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..ec856b5785 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' @@ -288,7 +289,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..93c51116b7 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/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 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, 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 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..40b2eae7d5 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. 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..7baa09c772 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, diff --git a/cuda-oxide/crates/cuda-device/src/shared.rs b/cuda-oxide/crates/cuda-device/src/shared.rs index 790e24078d..a6a7ed0c66 100644 --- a/cuda-oxide/crates/cuda-device/src/shared.rs +++ b/cuda-oxide/crates/cuda-device/src/shared.rs @@ -599,6 +599,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 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/tma.rs b/cuda-oxide/crates/cuda-device/src/tma.rs index fdbc2fb0be..bf5d72bdef 100644 --- a/cuda-oxide/crates/cuda-device/src/tma.rs +++ b/cuda-oxide/crates/cuda-device/src/tma.rs @@ -127,8 +127,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 +145,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..67395db9c7 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, @@ -1874,6 +1876,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 +1898,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 +1918,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 +1930,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/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/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/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/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..0adef5a939 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/epilogue/sm120.rs @@ -0,0 +1,547 @@ +/* + * 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 three = emit_i64_constant(session, &mut operations, 3, &input.location)?; + let warp_m = emit_binary( + session, + &mut operations, + "arith.andi", + input.operands[1].clone(), + three, + &input.location, + )?; + let two = emit_i64_constant(session, &mut operations, 2, &input.location)?; + let warp_n = emit_binary( + session, + &mut operations, + "arith.shrui", + input.operands[1].clone(), + two, + &input.location, + )?; + let lane = cast_int( + session, + &mut operations, + input.operands[2].clone(), + MlirType::Integer(64), + "arith.extui", + &input.location, + )?; + let fifteen = emit_i64_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_i64_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_i64_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_i64_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_i64_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_i64_constant(session, &mut operations, 7, &input.location)?; + let lane7 = emit_binary( + session, + &mut operations, + "arith.andi", + lane, + seven, + &input.location, + )?; + let eight = emit_i64_constant(session, &mut operations, 8, &input.location)?; + let swizzle = emit_binary( + session, + &mut operations, + "arith.muli", + lane7, + eight, + &input.location, + )?; + let half_stride = emit_i64_constant(session, &mut operations, 8192, &input.location)?; + + /// Logical column offsets for the eight accumulator slots within each M band. + const N_OFFSETS: [u64; 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_i64_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..17e2fae3e8 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/cute/mma/sm120.rs @@ -0,0 +1,1097 @@ +/* + * 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 = cast_int( + session, + &mut operations, + input.operands[0].clone(), + MlirType::Integer(64), + "arith.extui", + &input.location, + )?; + let four = emit_i64_constant(session, &mut operations, 4, &input.location)?; + let q = emit_binary( + session, + &mut operations, + "arith.divui", + lane.clone(), + four, + &input.location, + )?; + let three = emit_i64_constant(session, &mut operations, 3, &input.location)?; + let r = emit_binary( + session, + &mut operations, + "arith.andi", + lane, + three, + &input.location, + )?; + let one = emit_i64_constant(session, &mut operations, 1, &input.location)?; + let parity = emit_binary( + session, + &mut operations, + "arith.andi", + r, + one, + &input.location, + )?; + let eight = emit_i64_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_i64_constant(session, &mut operations, 16, &input.location)?; + let a_band = emit_binary( + session, + &mut operations, + "arith.muli", + input.operands[5].clone(), + sixteen.clone(), + &input.location, + )?; + let a_row0 = emit_binary( + session, + &mut operations, + "arith.addi", + a_band, + a_provider, + &input.location, + )?; + let sixty_four = emit_i64_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", + input.operands[6].clone(), + 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_u64, 32, 40, 64, 72, 96, 104] { + let delta = emit_i64_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 lane = cast_int( + session, + operations, + lane, + MlirType::Integer(64), + "arith.extui", + location, + )?; + let eight = emit_i64_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_i64_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_i64_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_i64_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_i64_constant(session, operations, LDMATRIX_BYTE_SWIZZLE_MASK, location)?; + let source_bits = emit_binary( + session, + operations, + "arith.andi", + plain.clone(), + mask, + location, + )?; + let three = emit_i64_constant(session, operations, LDMATRIX_BYTE_SWIZZLE_SHIFT, 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_i64_constant(session, operations, 31, location)?; + let low = emit_binary( + session, + operations, + "arith.andi", + row.clone(), + mask31, + location, + )?; + let four = emit_i64_constant(session, operations, 4, location)?; + let low = emit_binary(session, operations, "arith.muli", low, four, location)?; + let five = emit_i64_constant(session, operations, 5, location)?; + let quadrant = emit_binary(session, operations, "arith.shrui", row, five, location)?; + let three = emit_i64_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..384155c5ac --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/core.rs @@ -0,0 +1,1956 @@ +/* + * 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, + }, + 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_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..e03fb1de18 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/mir/memory.rs @@ -0,0 +1,1984 @@ +/* + * 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 pointee = pointer_pointee(ctx, source.deref(ctx).get_operand(0).get_type(ctx))?; + let target_pointee = session + .translate_type(pointee) + .map_err(|error| error.to_string())?; + let inbounds = MirPtrOffsetOp::new(source).is_inbounds(ctx); + let mut target = renamed("llvm.getelementptr", input)?; + gep_properties( + &mut target, + target_pointee, + vec![i32::MIN], + if inbounds { 3 } else { 0 }, + ); + Ok(vec![target]) + } +} + +/// 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) + ); + } +} 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..949e28dce9 --- /dev/null +++ b/cuda-oxide/crates/cuda-oxide-mlir-export/src/nvvm/sm120.rs @@ -0,0 +1,507 @@ +/* + * 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_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) + } +} + +/// 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-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/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/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..a6b24ac367 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/facts.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/facts.rs @@ -505,6 +505,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 +543,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 +591,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 +1082,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 +1090,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/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..7e6c95f899 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 @@ -572,11 +572,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..b707860ed2 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 @@ -81,6 +81,11 @@ pub fn translate_place( 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( 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..524771dea8 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. @@ -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..e3ad890548 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/types.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/types.rs @@ -670,6 +670,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(); @@ -764,7 +787,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 +902,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..35f975f1a0 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; 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/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/ops/cast.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs index 3b3b5b799f..6d9aff6b90 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs @@ -1743,6 +1743,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 +1923,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); 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/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/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.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.lock index c46db5e9e1..8f437f3864 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_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_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..6909f85289 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 @@ -105,6 +105,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 +136,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/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/src/main.rs index 4f1e2fd470..8511930604 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 @@ -22,7 +22,7 @@ //! - 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, @@ -88,6 +88,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 +104,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 +197,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 +208,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 // ============================================================================= @@ -905,8 +912,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)?; @@ -953,7 +964,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 +982,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 +1028,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"); 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/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/src/main.rs index 6b25119cb0..0dd24011ab 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,7 +7,7 @@ //! `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}; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock index 0df03ea524..66a470db38 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/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_api/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/src/main.rs index 12ba608b2b..8378f0c4f2 100644 --- 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 @@ -30,7 +30,7 @@ //! Usage: //! cargo oxide run shared_ptr_api -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::{DynamicSharedArray, SharedArray, SharedPtr, cuda_module, kernel, thread}; 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/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/src/main.rs index aa21c06446..e0028246fa 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 @@ -22,7 +22,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::{ 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 index acbd534d5a..f90ac2e9b2 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/images/tma-cta-local.svg +++ b/cuda-oxide/cuda-oxide-book/advanced/images/tma-cta-local.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/appendix/api-quick-reference.md b/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md index cc90fe89a5..b12efbf72a 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md +++ b/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md @@ -524,7 +524,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 +545,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. 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/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..c4fed98695 --- /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( + 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/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.lock new file mode 100644 index 0000000000..e45a7d5d5b --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.lock @@ -0,0 +1,794 @@ +# 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_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" +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-core", + "cuda-device", + "thiserror", +] + +[[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 = "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_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() as usize; + let warp = tid / 32; + let lane = tid % 32; + let base = DynamicSharedArray::::get_raw(); + unsafe { + let full = base.add(BARRIERS).cast::(); + let empty = full.add(STAGES); + if tid == 0 { + let mut s = 0; + while s < STAGES { + ManagedBarrier::::from_static(full.add(s)) + .init_instruction(1); + ManagedBarrier::::from_static(empty.add(s)) + .init_instruction(8); + s += 1; + } + } + fence::fence_mbarrier_init(); + barrier::barrier_cta_sync(0, THREADS as u32); + let mut tile = thread::blockIdx_x() as usize; + let stride = thread::gridDim_x() as usize; + let mut slot = 0usize; + if warp < 8 { + let wm = warp % 4; + let wn = warp / 4; + let mut phase = 0u32; + while tile < TOTAL_TILES { + let tile_m = tile % (M / 128); + let tile_n = tile / (M / 128); + let mut accumulators = AccumulatorTile16x8::<2, 8>::zero(); + let mut stage = 0; + while stage < K_STAGES { + let ready = + ManagedBarrier::::adopt_ready(full.add(slot)); + wait(&ready, phase); + let a = base.add(slot * DATA_BYTES); + let b = base.add(B_OFFSET + slot * DATA_BYTES); + let sa = base.add(SFA_OFFSET + slot * SCALE_BYTES).cast::(); + let sb = base.add(SFB_OFFSET + slot * 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 = 0; + while half < 2 { + let a0 = mma::load_mxf4_a(SharedAddress::from_raw( + a.add(physical::matrix_row(wm, half, lane)).cast(), + )); + let a1 = mma::load_mxf4_a(SharedAddress::from_raw( + a.add(physical::matrix_row(wm + 4, half, 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, lane)).cast(), + )); + accumulators.mma_mxf4_pair::( + [a0, a1], + [b0, b1], + [scale_a0, scale_a1], + scales_b.pair::(half), + ); + }); + half += 1; + } + if lane == 0 { + mbarrier::arrive( + &ManagedBarrier::::adopt_ready( + empty.add(slot), + ), + ); + } + slot += 1; + if slot == STAGES { + 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 { + let tile_m = tile % (M / 128); + let tile_n = tile / (M / 128); + let mut stage = 0; + while stage < K_STAGES { + wait( + &ManagedBarrier::::adopt_ready(empty.add(slot)), + phase, + ); + let ready = + ManagedBarrier::::adopt_ready(full.add(slot)); + mbarrier::arrive_expect_tx(&ready, TX_BYTES); + let a = region::(base.add(slot * DATA_BYTES)); + let b = region::(base.add(B_OFFSET + slot * DATA_BYTES)); + let sa = region::(base.add(SFA_OFFSET + slot * SCALE_BYTES)); + let sb = region::(base.add(SFB_OFFSET + slot * 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 + stage) as i32), + ready.shared_address(), + ); + tma::cp_async_bulk_tensor_shared_cta_global( + &sb, + sb_map, + Coordinates2D::new(0, (tile_n * K_STAGES + stage) as i32), + ready.shared_address(), + ); + slot += 1; + if slot == STAGES { + slot = 0; + phase ^= 1; + } + stage += 1; + } + tile += stride; + } + let mut tail = 0; + while tail < STAGES { + wait( + &ManagedBarrier::::adopt_ready(empty.add(slot)), + phase, + ); + slot += 1; + if slot == STAGES { + 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..edfc235f7a --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/src/main.rs @@ -0,0 +1,726 @@ +/* + * 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( + a_tma: *const TensorMap, + b_tma: *const TensorMap, + sfa_tma: *const TensorMap, + sfb_tma: *const TensorMap, + out_tma: *const 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 a_desc_dev = cuda_core::DeviceBuffer::from_host(&stream, &a_desc.bytes).unwrap(); + let b_desc_dev = cuda_core::DeviceBuffer::from_host(&stream, &b_desc.bytes).unwrap(); + let sfa_desc_dev = cuda_core::DeviceBuffer::from_host(&stream, &sfa_desc.bytes).unwrap(); + let sfb_desc_dev = cuda_core::DeviceBuffer::from_host(&stream, &sfb_desc.bytes).unwrap(); + let a_desc_ptr = a_desc_dev.cu_deviceptr() as *const TensorMap; + let b_desc_ptr = b_desc_dev.cu_deviceptr() as *const TensorMap; + let sfa_desc_ptr = sfa_desc_dev.cu_deviceptr() as *const TensorMap; + let sfb_desc_ptr = sfb_desc_dev.cu_deviceptr() as *const TensorMap; + + 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"); + let out_desc_dev = cuda_core::DeviceBuffer::from_host(&stream, &out_desc.bytes).unwrap(); + let out_desc_ptr = + out_desc_dev.cu_deviceptr() as *const TensorMap; + // 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. + module + .blockscale_gemm( + &stream, + &launch, + a_desc_ptr, + b_desc_ptr, + sfa_desc_ptr, + sfb_desc_ptr, + out_desc_ptr, + ) + .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 { + module + .blockscale_gemm( + &stream, + &launch, + a_desc_ptr, + b_desc_ptr, + sfa_desc_ptr, + sfb_desc_ptr, + out_desc_ptr, + ) + .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"); + module + .blockscale_gemm( + &stream, + &launch, + a_desc_ptr, + b_desc_ptr, + sfa_desc_ptr, + sfb_desc_ptr, + out_desc_ptr, + ) + .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/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.lock b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.lock new file mode 100644 index 0000000000..f88a780692 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.lock @@ -0,0 +1,794 @@ +# 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-core", + "cuda-device", + "thiserror", +] + +[[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 = "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 = "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" +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/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::