diff --git a/.github/workflows/book.yml b/.github/workflows/book.yml index b99d7a939a..732ff8b667 100644 --- a/.github/workflows/book.yml +++ b/.github/workflows/book.yml @@ -35,6 +35,8 @@ permissions: # Scope this to the ref, and keep the single-deploy guarantee on the deploy job # itself. Superseded PR runs are cancelled; pushes to main are not, so a # publish in flight is never interrupted. +# The group deliberately has no commit SHA. Parallel builds on main would let +# a slower older commit deploy after a newer one and replace the published site. concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: ${{ github.event_name == 'pull_request' }} diff --git a/.github/workflows/ci-status.yml b/.github/workflows/ci-status.yml new file mode 100644 index 0000000000..2c9f44b127 --- /dev/null +++ b/.github/workflows/ci-status.yml @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +--- +name: CI status + +# One status for the whole repository, so the README needs one badge. +# +# Badges are per workflow file, and the two pipelines have to stay separate +# for now: ci.yml (SIMT and the shared host crates) runs on GitHub-hosted +# runners on push and pull_request; cutile-rs.yml (Tile) runs on NVIDIA +# self-hosted runners, which cannot run on pull_request events, so it +# triggers on push to main and to the copy-pr-bot mirror branches. +# +# This workflow runs whenever either of them completes on main. For each +# pipeline it finds the newest completed push run on main that actually +# executed jobs and fails if that run failed. Two kinds of run say nothing +# about the pipeline and are skipped over: a run whose jobs were all skipped +# (cutile-rs.yml's `changes` gate does that on a non-Tile push) and a +# cancelled run. Taking either at face value would turn the badge green while +# the last real run was red. +# +# The badge in README.md points here. To cover another pipeline, add its +# workflow name to the trigger list and its file name to the loop below. +on: + workflow_run: + workflows: [CI, cutile-rs-pr] + types: [completed] + branches: [main] + +permissions: + actions: read + contents: read + +jobs: + status: + name: combined CI status on main + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - name: Report the newest executed main run of each pipeline + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + failed=0 + for workflow in ci.yml cutile-rs.yml; do + verdict=none + url=- + while read -r id conclusion run_url; do + [ -n "${id}" ] || continue + case "${conclusion}" in + success|failure|timed_out) ;; + *) continue ;; + esac + executed="$(gh api "repos/${GITHUB_REPOSITORY}/actions/runs/${id}/jobs?per_page=100" \ + --jq '[.jobs[] | select(.name != "changes" and .conclusion != "skipped")] | length')" + if [ "${executed}" -gt 0 ]; then + verdict="${conclusion}" + url="${run_url}" + break + fi + done < <(gh api \ + "repos/${GITHUB_REPOSITORY}/actions/workflows/${workflow}/runs?branch=main&event=push&status=completed&per_page=10" \ + --jq '.workflow_runs[] | "\(.id) \(.conclusion) \(.html_url)"') + echo "${workflow}: ${verdict} (${url})" + case "${verdict}" in + success|none) ;; + *) failed=1 ;; + esac + done + exit "${failed}" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 60f0f88fad..bd0c764954 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,9 +28,11 @@ on: - 'Cargo.lock' - '.github/workflows/**' +# 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. concurrency: - group: ci-gate-${{ github.ref_name }} - cancel-in-progress: true + group: ci-gate-${{ github.ref_name }}${{ github.ref_name == 'main' && format('-{0}', github.sha) || '' }} + cancel-in-progress: ${{ github.ref_name != 'main' }} # Every other workflow in the repo pins this rather than inheriting whatever the # repository default happens to be. This one is a pure caller today, so the four diff --git a/.github/workflows/clippy.yml b/.github/workflows/clippy.yml index 018779dbb9..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@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 b36896e9c3..60aaba1f3d 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -1,17 +1,36 @@ --- name: codeql +# Push and pull_request run only when Rust, workflow, or script inputs change. +# A Markdown-only commit skips both analyzers. The weekly schedule has no path +# filter, so that scan still covers the whole tree. on: push: branches: [main] + paths: + - '**/*.rs' + - '**/Cargo.toml' + - '**/Cargo.lock' + - '.github/**' + - '**/*.sh' + - '**/*.py' pull_request: branches: [main] + paths: + - '**/*.rs' + - '**/Cargo.toml' + - '**/Cargo.lock' + - '.github/**' + - '**/*.sh' + - '**/*.py' schedule: - cron: "44 3 * * 6" +# 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. concurrency: - group: ${{ github.workflow }}-${{ github.ref_name }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.ref_name }}${{ github.ref_name == 'main' && format('-{0}', github.sha) || '' }} + cancel-in-progress: ${{ github.ref_name != 'main' }} jobs: analyze: diff --git a/.github/workflows/cutile-rs.yml b/.github/workflows/cutile-rs.yml index 2c3735f6d7..11d0496217 100644 --- a/.github/workflows/cutile-rs.yml +++ b/.github/workflows/cutile-rs.yml @@ -25,9 +25,11 @@ on: - "pull-request/[0-9]+" workflow_dispatch: +# 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. concurrency: - group: ${{ github.workflow }}-${{ github.ref_name }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.ref_name }}${{ github.ref_name == 'main' && format('-{0}', github.sha) || '' }} + cancel-in-progress: ${{ github.ref_name != 'main' }} permissions: contents: read @@ -39,8 +41,11 @@ defaults: jobs: # Tile lane gate: cutile-rs tree plus shared host crates. Skips when only - # cuda-oxide changes. When copy-pr-bot yields no detectable file diff, run - # anyway so the mirror push is never silently skipped. + # cuda-oxide (or any other non-tile path) changes. `any_changed` is a real + # filter: dorny/paths-filter sets an output only for filters declared here, + # and an undeclared name stays empty, which would make `!= 'true'` always + # succeed. When copy-pr-bot yields no detectable file diff, both outputs are + # false and the jobs run anyway so the mirror push is never silently skipped. changes: runs-on: ubuntu-latest outputs: @@ -50,7 +55,7 @@ jobs: - uses: actions/checkout@v7 with: fetch-depth: 0 - - uses: dorny/paths-filter@v3 + - uses: dorny/paths-filter@0e4a8c6effa4802afeda77dc8d303f8176d7dfad # v3.0.4 id: changes with: filters: | @@ -65,6 +70,10 @@ jobs: - 'rust-toolchain.toml' - '.github/workflows/cutile-rs.yml' - '.github/copy-pr-bot.yaml' + # `*` covers a root file such as AGENTS.md. `**` covers every nested path. + any_changed: + - '*' + - '**' # Exercise the oldest Tile IR assembler, including its lack of # --list-versions. The separate 13.0 lane only checks shared host crates. @@ -199,7 +208,7 @@ jobs: # Prebuilt binary instead of `cargo install` (which compiled the tool # from source on every run). - - uses: taiki-e/install-action@9114bf4d891761788c546334fd37538eae1bf8b3 # v2.87.16 + - uses: taiki-e/install-action@9114bf4d891761788c546334fd37538eae1bf8b3 # v2.87.16 with: tool: cargo-semver-checks diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 5fd315f684..c12004ba20 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -26,9 +26,11 @@ on: - 'Cargo.lock' - '.github/workflows/**' +# 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. concurrency: - group: ${{ github.workflow }}-${{ github.ref_name }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.ref_name }}${{ github.ref_name == 'main' && format('-{0}', github.sha) || '' }} + cancel-in-progress: ${{ github.ref_name != 'main' }} permissions: contents: read @@ -48,10 +50,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@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 2b4ff38602..e481ee6e41 100644 --- a/.github/workflows/examples-compile.yml +++ b/.github/workflows/examples-compile.yml @@ -26,9 +26,11 @@ on: - 'Cargo.lock' - '.github/workflows/**' +# 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. concurrency: - group: ${{ github.workflow }}-${{ github.ref_name }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.ref_name }}${{ github.ref_name == 'main' && format('-{0}', github.sha) || '' }} + cancel-in-progress: ${{ github.ref_name != 'main' }} permissions: contents: read @@ -59,10 +61,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@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 @@ -74,30 +76,31 @@ jobs: # The device pipeline shells out to LLVM's llc/opt for PTX # generation. This pin is the FLOOR GATE: if someone lands IR that # needs a newer llc, this job is where the floor bump becomes - # visible instead of on contributors' machines. LLVM 22 is the true - # floor: the ldmatrix/stmatrix family (llvm.nvvm.stmatrix.*, used by - # gemm_sol and gemm_sol_final) postdates llc-21, which "lowers" an - # unknown intrinsic by leaking its dotted LLVM name into the PTX as - # an `.extern .func` (exit 0, no diagnostic): unassemblable output - # that only surfaced once the ptxas verdict gate started assembling - # what this job produces. The pipeline now rejects such leaks with a - # named error, so an under-floor llc fails loudly here. - # If a future ubuntu-latest codename has no llvm-toolchain-*-22 + # visible instead of on contributors' machines. LLVM 23 is the true + # floor: the bf16x2 to packed-FP8 conversions (PTX 9.1 forms, used + # by generated_intrinsics_blackwell) need `+ptx91`, which llc-22 + # does not know. It ignored the unknown feature and emitted + # `.version 8.8` PTX that failed only at JIT; the pipeline now + # refuses a PTX 9.1 request below LLVM 23 with a named error, so an + # under-floor llc fails loudly here. The previous floor, LLVM 22 for + # the ldmatrix/stmatrix family (llc-21 leaked unknown intrinsics + # into the PTX as `.extern .func`), is covered by the same gate. + # If a future ubuntu-latest codename has no llvm-toolchain-*-23 # suite, apt-get update fails loudly; recover by pinning runs-on to # the previous image or bumping the LLVM major here. - - name: Install LLVM 22 (llc/opt for the NVPTX backend) + - name: Install LLVM 23 (llc/opt for the NVPTX backend) 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}-22 main" \ - | sudo tee /etc/apt/sources.list.d/llvm-22.list + echo "deb [signed-by=/etc/apt/keyrings/apt-llvm-org.asc] https://apt.llvm.org/${VERSION_CODENAME}/ llvm-toolchain-${VERSION_CODENAME}-23 main" \ + | sudo tee /etc/apt/sources.list.d/llvm-23.list sudo apt-get update - sudo apt-get install -y --no-install-recommends llvm-22 - llc-22 --version | head -2 - opt-22 --version | head -2 + sudo apt-get install -y --no-install-recommends llvm-23 + llc-23 --version | head -2 + opt-23 --version | head -2 - name: Compile every example (smoketest --compile-only) # Step-level timeout, below the job's 75: a job-level timeout @@ -106,7 +109,7 @@ jobs: timeout-minutes: 50 env: CUDA_TOOLKIT_PATH: ${{ steps.cuda-toolkit.outputs.CUDA_PATH }} - CUDA_OXIDE_LLC: /usr/bin/llc-22 + CUDA_OXIDE_LLC: /usr/bin/llc-23 # Each example is its own cargo workspace; without a shared # target dir the dep tree rebuilds ~85 times and blows the # runner's ~14 GB disk (same trick as clippy.yml). smoketest.sh @@ -153,7 +156,7 @@ jobs: working-directory: cuda-oxide env: CUDA_TOOLKIT_PATH: ${{ steps.cuda-toolkit.outputs.CUDA_PATH }} - CUDA_OXIDE_LLC: /usr/bin/llc-22 + CUDA_OXIDE_LLC: /usr/bin/llc-23 CARGO_TARGET_DIR: ${{ runner.temp }}/examples-shared-target CARGO_INCREMENTAL: "0" run: | 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 5a5a09b6fd..00eb30f37e 100644 --- a/.github/workflows/naming-guard.yml +++ b/.github/workflows/naming-guard.yml @@ -1,15 +1,32 @@ --- name: naming-guard +# check-reserved-prefixes.sh searches only cuda-oxide/crates. The strings it +# looks for are the SIMT link-symbol prefixes owned by reserved-oxide-symbols; +# every other SIMT crate must use that crate's constants. cutile-rs and the +# shared host crates (cuda-bindings, cuda-core, cuda-core-derive, cuda-async) +# are outside that rule, so these paths omit them. Listing those directories +# here would rerun the same search; covering them means adding them to +# SEARCH_ROOTS in the script first. on: push: branches: [main] + paths: + - 'cuda-oxide/crates/**' + - 'cuda-oxide/scripts/check-reserved-prefixes.sh' + - '.github/workflows/naming-guard.yml' pull_request: branches: [main] + paths: + - 'cuda-oxide/crates/**' + - 'cuda-oxide/scripts/check-reserved-prefixes.sh' + - '.github/workflows/naming-guard.yml' +# 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. concurrency: - group: ${{ github.workflow }}-${{ github.ref_name }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.ref_name }}${{ github.ref_name == 'main' && format('-{0}', github.sha) || '' }} + cancel-in-progress: ${{ github.ref_name != 'main' }} permissions: contents: read diff --git a/.github/workflows/status-guard.yml b/.github/workflows/status-guard.yml index b56320c346..56deee2e06 100644 --- a/.github/workflows/status-guard.yml +++ b/.github/workflows/status-guard.yml @@ -7,9 +7,11 @@ on: pull_request: branches: [main] +# 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. concurrency: - group: ${{ github.workflow }}-${{ github.ref_name }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.ref_name }}${{ github.ref_name == 'main' && format('-{0}', github.sha) || '' }} + cancel-in-progress: ${{ github.ref_name != 'main' }} permissions: contents: read @@ -58,13 +60,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 @@ -128,9 +129,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 @@ -228,11 +229,12 @@ jobs: - name: Verify docs and templates spell the shared crates' simt paths run: bash cuda-oxide/scripts/check-host-api-paths.sh - # The shared host-crate release is named once in the root workspace and - # copied into every example manifest and into the `cargo oxide new` - # templates. A drifted example resolves a second runtime and breaks the - # shared example build cache; a drifted scaffold hands each new project a - # runtime the generated code was not written against. + # The shared host crates are path dependencies in the SIMT and example + # workspaces, and the `cargo oxide new` templates take them from the same + # git source as cuda-host. A drifted example resolves a second runtime and + # breaks the shared example build cache; a scaffold naming a crates.io copy + # hands each new project two cuda_core crates, and the project fails to + # compile. shared-crate-pin-parity: name: shared host-crate pin parity runs-on: ubuntu-latest @@ -242,7 +244,7 @@ jobs: uses: actions/checkout@v4 # Reads sources only: no rustup, no cargo, no CUDA toolkit. - - name: Verify every shared-crate pin agrees with the root workspace + - name: Verify the host crates are path dependencies at tip of tree and the scaffold takes them from git beside cuda-host run: bash cuda-oxide/scripts/check-shared-crate-pin.sh # cuda-oxide/crates/oxide-artifacts is only the publish source now; the diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index cf9a33179b..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 @@ -68,6 +80,11 @@ jobs: # tests that cover the configuration actually built. - package: oxide-artifacts test_args: --all-targets --features object + # `#[tile_export]` descriptions and their validation rules, shared + # by cuda-macros, the backend and cargo-oxide. Pure Rust, no CUDA. + - package: oxide-tile-export + # The Tile IR bytecode writer for those descriptions. Pure Rust. + - package: oxide-tilebc - package: reserved-oxide-symbols # launch_requires_unsafe trybuild fixture dev-depends on # cuda-host/cuda-core, which transitively pull cuda-bindings @@ -97,29 +114,14 @@ jobs: - name: Install CUDA toolkit if: matrix.needs_cuda - uses: Jimver/cuda-toolkit@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/.vscode/settings.json b/.vscode/settings.json index f04c586f3e..a06fba831c 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -2,7 +2,7 @@ // Project-level editor settings, tracked in git. Keep personal preferences // (fonts, colours, keybindings) in your VS Code User settings instead. "rust-analyzer.rustc.source": "discover", - "rust-analyzer.cargo.features": "all", + "rust-analyzer.cargo.features": [], "rust-analyzer.procMacro.enable": true, "rust-analyzer.procMacro.attributes.enable": true, // "rust-analyzer.cargo.extraEnv": { diff --git a/AGENTS.md b/AGENTS.md index 9024129256..85a13cdb67 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 @@ -67,8 +68,8 @@ all CI: example compilation, additional example lint scopes, book builds, and CodeQL have separate jobs. Read the recipe and relevant workflow when choosing validation for a change. -For Tile and shared-host changes, follow `cutile-rs/CONTRIBUTING.md` and -`.github/workflows/cutile-rs.yml`. The CPU and assembler suite in +For Tile and shared-host changes, follow the cutile-rs section of +`CONTRIBUTING.md` and `.github/workflows/cutile-rs.yml`. The CPU and assembler suite in `cutile-rs/scripts/run_cpu_tests.sh` runs with `cutile-rs/` as its working directory and includes the CPU-safe shared-host tests. CPU-only does not mean toolkit-free: some suites require CUDA headers and the Tile IR assembler. GPU tests have a 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. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0200b930a0..93c51116b7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -19,6 +19,7 @@ If you are unsure whether something is worth a full issue or PR, the Discord - [Signing Your Commits](#signing-your-commits) - [Contribution Process](#contribution-process) - [Code Requirements](#code-requirements) +- [AI-Assisted Contributions](#ai-assisted-contributions) - [IP Review Process](#ip-review-process) ## Developer Certificate of Origin @@ -111,22 +112,46 @@ Your `Signed-off-by` name and email must match your Git configuration Pull requests that do not meet the requirements below or lack proper DCO sign-off will not be merged. +### Branch and pull request naming + +Prefix branches and pull request titles with the kind of change. Branches use +`/` as the separator and titles use `:`, lowercase: `fix/kernel-launch-sync` +and `fix: resolve kernel launch sync race`. + +| Prefix | Use for | +| --- | --- | +| `feat` | New features | +| `fix` | Bug fixes | +| `docs` | Documentation changes | +| `refactor` | Code changes with no behavior change | +| `perf` | Performance improvements with no functional change | +| `test` | Test additions or changes | +| `build` | Build system and dependency changes | +| `ci` | CI changes | +| `chore` | Maintenance tasks | + +A scope in parentheses is welcome when a change is confined to one component, +for example `docs(cutile-rs): ...`. + ## Code Requirements ### Toolchains The shared host crates at the repository root use the root stable pin. Each product has its own workspace and toolchain under `cuda-oxide/` or -`cutile-rs/`. For Tile development, see [cutile-rs/CONTRIBUTING.md](cutile-rs/CONTRIBUTING.md). -The setup and validation commands below cover cuda-oxide. +`cutile-rs/`. The setup commands below cover cuda-oxide; the Tile checks have +their own section further down. -cuda-oxide requires the Rust nightly toolchain with `rustc_private` support. -See the [cuda-oxide README](cuda-oxide/README.md) for setup instructions. +cuda-oxide pins its stable Rust toolchain (`1.99.1`) and required `rustc-dev` +and `llvm-tools` components in `cuda-oxide/rust-toolchain.toml` (with +`RUSTC_BOOTSTRAP=1` configured via `cuda-oxide/.cargo/config.toml` and +`cargo-oxide`). See the [cuda-oxide README](cuda-oxide/README.md) for setup +instructions. -`cuda-oxide/flake.nix` provides the CUDA 13, LLVM 22, Clang, and pinned nightly -development environment. From the repository root, run `cd cuda-oxide` and -then `nix develop` with flakes enabled. Container setup is documented in -[cuda-oxide/.devcontainer/README.md](cuda-oxide/.devcontainer/README.md). +`cuda-oxide/flake.nix` provides the CUDA 13, Clang/libclang, and pinned Rust +(with LLVM 23 `llvm-tools`) development environment. From the repository root, +run `cd cuda-oxide` and then `nix develop` with flakes enabled. Container setup +is documented in [cuda-oxide/.devcontainer/README.md](cuda-oxide/.devcontainer/README.md). ### Running the cuda-oxide checks @@ -150,6 +175,28 @@ codegen backend, a Python virtualenv, or GitHub's own infrastructure. The step when a workflow changes; the commands below are the ones worth knowing by hand even so. +### Running the cutile-rs checks + +cutile-rs builds on the stable toolchain pinned in `cutile-rs/rust-toolchain.toml` +and needs a CUDA 13.2 or newer toolkit with the Tile IR assembler, `tileiras`; +the [cutile-rs book](https://nvidia.github.io/cuda-rust/cutile/) covers setup. +Its test scripts locate the Tile workspace from their own path and reach the +shared host crates at the repository root on their own, so they can be run +from anywhere: + +```bash +cutile-rs/scripts/run_cpu_tests.sh # CPU and assembler tests; needs the toolkit, not a GPU +cutile-rs/scripts/run_gpu_tests.sh # GPU integration tests +cutile-rs/scripts/run_all.sh # both of the above, then the examples and benchmarks +``` + +Run `cutile-rs/scripts/run_all.sh` before marking a Tile pull request ready +for review. +CI runs the same suites in `.github/workflows/cutile-rs.yml`. Because the Tile +lanes use NVIDIA self-hosted runners, that workflow triggers on the +`pull-request/` branches that copy-pr-bot mirrors rather than on the pull +request itself. + ### Formatting and Style - Run `(cd cuda-oxide && cargo oxide fmt)` before submitting. All code must be @@ -247,6 +294,14 @@ error's `Display` names the library candidates the loader tried. files directly and asks for a row per third-party crate. Examples that depend only on first-party crates by path need nothing. +## AI-Assisted Contributions + +- **You are responsible:** Review and understand all AI-generated code before submitting it. +- **Verify correctness:** AI-generated code may contain bugs or security issues. Run the applicable tests and review the diff carefully. +- **No AI-generated sensitive content:** Do not use AI to generate security policies, legal, or license text. + +See [AGENTS.md](AGENTS.md) for instructions for coding agents working in this repository. + ## IP Review Process All contributions to CUDA Rust are subject to NVIDIA's IP review process. diff --git a/Cargo.lock b/Cargo.lock index 427b364038..ca0b7d7ad7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -105,7 +105,7 @@ checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -121,7 +121,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -133,7 +133,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -145,7 +145,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index e27f86ecc4..efd77cfc1c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ 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", diff --git a/README.md b/README.md index 23ee8c5bdb..9fe71ec86b 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,10 @@ # CUDA Rust +

+ Build + Docs +

+ CUDA Rust is NVIDIA's CUDA platform for Rust. It provides host-side crates for working with the CUDA driver and device-side crates and tools for writing GPU kernels in idiomatic Rust. @@ -37,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 | @@ -76,19 +81,13 @@ 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 CUDA Rust is licensed under the [Apache License 2.0](./LICENSE). The root -`LICENSE` covers the whole repository. Each component also includes a copy with -its sources and built artifacts. - -| Component | License | License file | -| --- | --- | --- | -| `cuda-oxide` | Apache-2.0 | root [`LICENSE`](./LICENSE) | -| `cutile-rs` | Apache-2.0 | [`cutile-rs/LICENSE`](./cutile-rs/LICENSE) | -| `cuda-bindings`, `cuda-core`, `cuda-core-derive`, `cuda-async` | Apache-2.0 | root [`LICENSE`](./LICENSE) | +`LICENSE` covers the whole repository: cuda-oxide, cutile-rs, and the shared +host crates. Every published crate declares `license = "Apache-2.0"`. Third-party attributions are listed in [`cuda-oxide/THIRD_PARTY_NOTICES`](./cuda-oxide/THIRD_PARTY_NOTICES). diff --git a/cuda-async/Cargo.toml b/cuda-async/Cargo.toml index 29a2eef8cf..24fd68a5d9 100644 --- a/cuda-async/Cargo.toml +++ b/cuda-async/Cargo.toml @@ -1,19 +1,19 @@ [package] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" edition = "2021" rust-version = "1.89" license = "Apache-2.0" authors = ["Melih Elibol"] -repository = "https://github.com/nvlabs/cutile-rs" +repository = "https://github.com/NVIDIA/cuda-rust" readme = "README.md" 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-bindings/Cargo.toml b/cuda-bindings/Cargo.toml index 473e1fae67..eee40b94ed 100644 --- a/cuda-bindings/Cargo.toml +++ b/cuda-bindings/Cargo.toml @@ -1,13 +1,13 @@ [package] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" authors = ["Melih Elibol"] edition = "2021" rust-version = "1.89" description = "NVIDIA CUDA bindings." readme = "README.md" license = "Apache-2.0" -repository = "https://github.com/nvlabs/cutile-rs" +repository = "https://github.com/NVIDIA/cuda-rust" [lib] doctest = false diff --git a/cuda-core-derive/Cargo.toml b/cuda-core-derive/Cargo.toml index 7bebcc7c3d..797531a8b5 100644 --- a/cuda-core-derive/Cargo.toml +++ b/cuda-core-derive/Cargo.toml @@ -1,12 +1,12 @@ [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" authors = ["Melih Elibol"] -repository = "https://github.com/nvlabs/cutile-rs" +repository = "https://github.com/NVIDIA/cuda-rust" [lib] proc-macro = true diff --git a/cuda-core/Cargo.toml b/cuda-core/Cargo.toml index 2ef3b79c76..98bf503077 100644 --- a/cuda-core/Cargo.toml +++ b/cuda-core/Cargo.toml @@ -1,20 +1,20 @@ [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" authors = ["Melih Elibol"] -repository = "https://github.com/nvlabs/cutile-rs" +repository = "https://github.com/NVIDIA/cuda-rust" readme = "README.md" [dependencies] # 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" } anyhow = "1" half = { version = "2" } oxide-artifacts = { version = "0.2.1", features = ["object-read"] } 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 26f3fbd4a5..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-08-28` 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 577857720e..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-08-28", + "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 5bb9ee62fa..adf2d64339 100644 --- a/cuda-oxide/Cargo.lock +++ b/cuda-oxide/Cargo.lock @@ -81,9 +81,9 @@ checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" [[package]] name = "awint" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d48b7360a36d335663e8f20b7f439029debf52b5a0a749d6e936405f25045288" +checksum = "b2a8f298efd080b0d3e1ee1723b5d88ecc0d78805fad347e09b76dbd4dadbe37" dependencies = [ "awint_core", "awint_dag", @@ -94,9 +94,9 @@ dependencies = [ [[package]] name = "awint_core" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8280079e78217ace721501f35dbf551dadf0ab63863504169316cea1bbac872" +checksum = "c24d226d9faa090abec192c8ad1ea74635e6cc94000c72a95ceb60cf691428c3" dependencies = [ "awint_internals", "const_fn", @@ -104,9 +104,9 @@ dependencies = [ [[package]] name = "awint_dag" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73977ba1786a5e127272f08b865b60f5d548576a0925de1daa5703cc9ba78a8d" +checksum = "dc5377016f3077b18304da4233a28626ea069047ca84c9bc8c8b81b659f11032" dependencies = [ "awint_ext", "awint_macro_internals", @@ -116,9 +116,9 @@ dependencies = [ [[package]] name = "awint_ext" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef1b30d3640d9f94dd9f55f655773c5238c11be6f95ab36d24308b7ae34c9bbc" +checksum = "01ff843cdefdd48fbe9bc11292a902b7921ff37433ed4124fa02136f83186f58" dependencies = [ "awint_core", "const_fn", @@ -126,18 +126,18 @@ dependencies = [ [[package]] name = "awint_internals" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "173937e3f4e233d93362fc1e28ab23808cb31671a1923033706cb44ed4e5d10c" +checksum = "6209d7873685ac95ee41c4b63fe29bff9b2e1106de65d05b106b2b769c178575" dependencies = [ "const_fn", ] [[package]] name = "awint_macro_internals" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2c6a23f64f14c3b566c2c04124015f78c5fb93e8275a6574e0118747ee60de6" +checksum = "6b519a9cd37bf73e44e29b3779ca90246412b56db4af4e96d959f5082bb8bdcc" dependencies = [ "awint_ext", "proc-macro2", @@ -146,11 +146,10 @@ dependencies = [ [[package]] name = "awint_macros" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16e50a2862c6002a424fc9bb4982ec33ab679cb4ad04a174e2a52c2ecc3ea4c7" +checksum = "518a133c2163e8f021f708d7278af27dd6c1bcee96ea672a2d825b13a6931818" dependencies = [ - "awint_internals", "awint_macro_internals", ] @@ -179,9 +178,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.11.1" +version = "2.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" [[package]] name = "block-buffer" @@ -213,6 +212,8 @@ dependencies = [ "libnvvm-sys", "nvjitlink-sys", "oxide-artifacts 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", + "oxide-tile-export", + "oxide-tilebc", "ptx-parse", "ptx-schedule", "reserved-oxide-symbols", @@ -295,9 +296,9 @@ checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" [[package]] name = "combine" -version = "4.6.7" +version = "4.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" dependencies = [ "bytes", "memchr", @@ -311,9 +312,9 @@ checksum = "413d67b29ef1021b4d60f4aa1e925ca031751e213832b4b1d588fae623c05c60" [[package]] name = "convert_case" -version = "0.11.0" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "affbf0190ed2caf063e3def54ff444b449371d55c58e513a95ab98eca50adb49" +checksum = "1af709f1f33454bf52eadfc8c78b3b9ef9cb26fb54d16dc9cd9a7299f899fd1b" dependencies = [ "unicode-segmentation", ] @@ -358,6 +359,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" @@ -367,12 +379,12 @@ dependencies = [ "serde", "serde_json", "sha2", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -381,13 +393,13 @@ dependencies = [ "futures", "half", "once_cell", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "thiserror 1.0.69", ] [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -399,7 +411,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -410,7 +422,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -440,7 +452,7 @@ dependencies = [ "oxide-artifacts 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", "ptx-parse", "sha2", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] @@ -464,9 +476,11 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "ctm-rs", "cuda-core", "cuda-device", "cuda-host", + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -478,27 +492,68 @@ 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.18", + "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-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" @@ -513,6 +568,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" @@ -522,7 +591,7 @@ dependencies = [ "linkme", "pliron", "pliron-derive", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] @@ -534,7 +603,7 @@ dependencies = [ "once_cell", "pliron", "pliron-derive", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] @@ -547,7 +616,7 @@ dependencies = [ "linkme", "pliron", "pliron-derive", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] @@ -801,9 +870,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hi_sparse_bitset" -version = "0.8.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad02fc3cac722c3c44c3ed708f6653e7d149b7912f7342fc1e8530455f84f0eb" +checksum = "682097abdf22a902b7b2b5f2824e773691507289873c3ea0abbc5d4c85f05ec0" dependencies = [ "wide", ] @@ -823,14 +892,14 @@ version = "0.1.0" dependencies = [ "dialect-iket", "pliron", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] name = "indexmap" -version = "2.14.0" +version = "2.14.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" dependencies = [ "equivalent", "hashbrown 0.17.1", @@ -838,9 +907,9 @@ dependencies = [ [[package]] name = "inventory" -version = "0.3.24" +version = "0.3.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4f0c30c76f2f4ccee3fe55a2435f691ca00c0e4bd87abe4f4a851b1d4dac39b" +checksum = "6928282826c822ad91bf1c9a1cb90a30ba1c26770749929b4656cd6be829cd7c" dependencies = [ "rustversion", ] @@ -900,27 +969,27 @@ version = "0.2.1" dependencies = [ "cuda-target-spec", "libloading", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] name = "linkme" -version = "0.3.36" +version = "0.3.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e83272d46373fb8decca684579ac3e7c8f3d71d4cc3aa693df8759e260ae41cf" +checksum = "3045e122bd98aef8ec3ad58ce84f0791f64e70163d1a02710af4aa11a4d54cc5" dependencies = [ "linkme-impl", ] [[package]] name = "linkme-impl" -version = "0.3.36" +version = "0.3.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32d59e20403c7d08fe62b4376edfe5c7fb2ef1e6b1465379686d0f21c8df444b" +checksum = "77060ebe535362c3da75682cd17b0431017b6e7c5661e714fc69a7ad017d1301" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.9", ] [[package]] @@ -950,8 +1019,8 @@ dependencies = [ "pliron-derive", "pliron-llvm", "reserved-oxide-symbols", - "rustc-hash 2.1.2", - "thiserror 2.0.18", + "rustc-hash 2.1.3", + "thiserror 2.0.21", ] [[package]] @@ -965,9 +1034,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.29" +version = "0.4.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" [[package]] name = "memchr" @@ -987,6 +1056,8 @@ version = "0.2.1" dependencies = [ "combine", "cuda-oxide-codegen", + "cute-layout", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", @@ -998,8 +1069,8 @@ dependencies = [ "pliron", "pliron-derive", "reserved-oxide-symbols", - "rustc-hash 2.1.2", - "thiserror 2.0.18", + "rustc-hash 2.1.3", + "thiserror 2.0.21", ] [[package]] @@ -1015,8 +1086,8 @@ dependencies = [ "pliron", "pliron-derive", "reserved-oxide-symbols", - "rustc-hash 2.1.2", - "thiserror 2.0.18", + "rustc-hash 2.1.3", + "thiserror 2.0.21", ] [[package]] @@ -1026,8 +1097,8 @@ dependencies = [ "dialect-mir", "pliron", "reserved-oxide-symbols", - "rustc-hash 2.1.2", - "thiserror 2.0.18", + "rustc-hash 2.1.3", + "thiserror 2.0.21", ] [[package]] @@ -1045,7 +1116,7 @@ name = "nvjitlink-sys" version = "0.2.1" dependencies = [ "libloading", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] @@ -1082,9 +1153,9 @@ checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "once_vec" -version = "0.3.0" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a6298074a1742b584aaca148a6c61ecb27f0b55fdf237c235939c84924377a1" +checksum = "0b61fbf7ec37342ce171d9f5e3b82be74bda6c4e84e3b28338a03d553913e6aa" [[package]] name = "oxide-artifacts" @@ -1102,6 +1173,22 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.21", +] + +[[package]] +name = "oxide-tilebc" +version = "0.2.1" +dependencies = [ + "thiserror 2.0.21", +] + [[package]] name = "parking_lot_core" version = "0.9.12" @@ -1123,8 +1210,9 @@ checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" [[package]] name = "pliron" -version = "0.18.0" -source = "git+https://github.com/pliron-org/pliron?rev=06d99a3df938394583715aacf5eadb65b4791c09#06d99a3df938394583715aacf5eadb65b4791c09" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7395b93d3fc99ee1923027fedead143eb63e4c76b6f6f6d79ebfac6698956317" dependencies = [ "awint", "combine", @@ -1139,38 +1227,51 @@ dependencies = [ "once_vec", "pliron-derive", "regex", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "rustc_apfloat", "slotmap", "smallvec", "spin", - "thiserror 2.0.18", + "thiserror 2.0.21", "utf8-chars", ] [[package]] name = "pliron-derive" -version = "0.18.0" -source = "git+https://github.com/pliron-org/pliron?rev=06d99a3df938394583715aacf5eadb65b4791c09#06d99a3df938394583715aacf5eadb65b4791c09" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "245f8ea25d8ddcaf04e3342be709b4f257cef2c2c15b99d263955d6f073796cd" dependencies = [ "combine", "convert_case", "indexmap", "proc-macro2", "quote", - "rustc-hash 2.1.2", - "syn 3.0.3", + "rustc-hash 2.1.3", + "syn 3.0.9", ] [[package]] name = "pliron-llvm" -version = "0.18.0" -source = "git+https://github.com/pliron-org/pliron?rev=06d99a3df938394583715aacf5eadb65b4791c09#06d99a3df938394583715aacf5eadb65b4791c09" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b027d3174757ccc3a96ae971aa4e7904439a3082a3550bc01361369c3ab073f" dependencies = [ "bitflags", "log", "pliron", - "thiserror 2.0.18", + "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]] @@ -1185,9 +1286,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.106" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] @@ -1206,14 +1307,14 @@ dependencies = [ "ptx-parse", "serde", "serde_json", - "thiserror 2.0.18", + "thiserror 2.0.21", ] [[package]] name = "quote" -version = "1.0.45" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] @@ -1226,9 +1327,9 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "recasting" -version = "0.2.1" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70aeee1a07922bfe7c79e148553dee262248d2493d9778eee7551b9b5d7cc651" +checksum = "dc3f51d3a1399f13fa081d97345d0fdf08597fe78558fc8539d0e67412f83b26" [[package]] name = "redox_syscall" @@ -1241,9 +1342,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.12.3" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" dependencies = [ "aho-corasick", "memchr", @@ -1253,9 +1354,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.14" +version = "0.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" dependencies = [ "aho-corasick", "memchr", @@ -1264,9 +1365,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.10" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" @@ -1286,9 +1387,9 @@ checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" [[package]] name = "rustc-hash" -version = "2.1.2" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustc_apfloat" @@ -1334,9 +1435,9 @@ checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "safe_arch" -version = "0.7.4" +version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" +checksum = "42c6efa15875e6ecb39ca61fb0b0c1a40b84fac5a5ffe71eef7d1000c8eb3f5f" dependencies = [ "bytemuck", ] @@ -1433,15 +1534,15 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.15.1" +version = "1.16.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" [[package]] name = "spin" -version = "0.12.0" +version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1527984ca054dfca79333baec451042863f485fbee01b7bf6d911de915cac865" +checksum = "0134f9043ed38b087ac4f7d4af44c79e2c9e5094421fe3164f435ce585953b10" dependencies = [ "lock_api", ] @@ -1465,9 +1566,9 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.3" +version = "3.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" dependencies = [ "proc-macro2", "quote", @@ -1513,11 +1614,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" dependencies = [ - "thiserror-impl 2.0.18", + "thiserror-impl 2.0.21", ] [[package]] @@ -1533,13 +1634,13 @@ dependencies = [ [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.9", ] [[package]] @@ -1583,9 +1684,9 @@ checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" [[package]] name = "triple_arena" -version = "0.14.0" +version = "0.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a99051ec05383f70bcab71884d39083a455257795e71cb780048a4bc890f3b8" +checksum = "eee55b64228381514ddc4ad004e631ccf45e594e428d13ce197697ce56ef1c55" dependencies = [ "recasting", ] @@ -1658,9 +1759,9 @@ dependencies = [ [[package]] name = "wide" -version = "0.7.33" +version = "1.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce5da8ecb62bcd8ec8b7ea19f69a51275e91299be594ea5cc6ef7819e16cd03" +checksum = "d920ac99c3c8edce110cb8d07dbb324d6d026011dce85b1e9355b70f0adacc4f" dependencies = [ "bytemuck", "safe_arch", diff --git a/cuda-oxide/Cargo.toml b/cuda-oxide/Cargo.toml index 9f89792f94..ec0021d2eb 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", @@ -24,6 +30,8 @@ members = [ "crates/nvvm-transforms", "crates/cargo-oxide", "crates/oxide-artifacts", + "crates/oxide-tile-export", + "crates/oxide-tilebc", # FFI bindings "crates/libnvvm-sys", "crates/nvjitlink-sys", @@ -37,36 +45,44 @@ members = [ edition = "2024" authors = ["Nihal Pasham "] license = "Apache-2.0" -repository = "https://github.com/NVlabs/cuda-oxide.git" +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 upstream struct-layout, GEP no-wrap, volatile, -# global-constant, inline-asm, and generic LLVM attribute semantics. +# Pliron MLIR framework, pinned to the published 0.20.0 release, retaining +# struct-layout, GEP no-wrap, volatile, global-constant, inline-asm, and generic +# LLVM attribute semantics. # pliron-llvm uses default-features = false to keep llvm-sys (LLVM 23) out: # cuda-oxide consumes the dialect modeling only, not the FFI bridge. -# `crates/rustc-codegen-cuda/Cargo.toml` MUST stay on this same rev, or pliron -# resolves to two distinct crates and types stop unifying. -pliron = { git = "https://github.com/pliron-org/pliron", rev = "06d99a3df938394583715aacf5eadb65b4791c09" } -pliron-derive = { git = "https://github.com/pliron-org/pliron", rev = "06d99a3df938394583715aacf5eadb65b4791c09" } -pliron-llvm = { git = "https://github.com/pliron-org/pliron", rev = "06d99a3df938394583715aacf5eadb65b4791c09", default-features = false, features = [ +# `crates/rustc-codegen-cuda/Cargo.toml` MUST use the same registry source and +# exact release, or pliron resolves to distinct crates and types stop unifying. +pliron = { version = "=0.20.0" } +pliron-derive = { version = "=0.20.0" } +pliron-llvm = { version = "=0.20.0", default-features = false, features = [ "std", ] } # 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" } dialect-nvvm = { path = "crates/dialect-nvvm" } mir-lower = { path = "crates/mir-lower" } iket-lower = { path = "crates/iket-lower" } -mir-importer = { path = "crates/mir-importer" } cuda-oxide-codegen = { path = "crates/cuda-oxide-codegen" } mir-transforms = { path = "crates/mir-transforms" } nvvm-transforms = { path = "crates/nvvm-transforms" } @@ -78,19 +94,18 @@ cuda-host = { path = "crates/cuda-host" } # `cuda-host` names `features = ["host"]`, so anything that loads or launches a # module still gets it. cuda-macros = { path = "crates/cuda-macros", default-features = false } -# Git-root host crates. `=` is an exact pin, same as cutile-rs, so a 0.4.1 -# package version does not match until this line is updated with it. -# Examples path-depend on the same directories so the graph has one -# cuda_core (and one cuda-async / cuda-bindings). `cargo oxide new` still -# pins the crates.io version of this line for out-of-tree projects; keep -# SHARED_HOST_CRATES_VERSION in lockstep (the pin check ignores the `=`). -cuda-bindings = { path = "../cuda-bindings", version = "=0.4.0" } -cuda-core = { path = "../cuda-core", version = "=0.4.0" } -cuda-async = { path = "../cuda-async", version = "=0.4.0" } +# Git-root host crates, consumed by path at tip of tree. No version +# requirement on purpose: a release of the host crates needs no edit here, and +# the crates that depend on them in this tree are not published, so nothing +# needs one (#1459). Examples path-depend on the same directories so the graph +# has one cuda_core (and one cuda-async / cuda-bindings). `cargo oxide new` +# names every CUDA crate from one git source for the same reason; +# scripts/check-shared-crate-pin.sh enforces both. +cuda-core = { path = "../cuda-core" } +cuda-async = { path = "../cuda-async" } libnvvm-sys = { path = "crates/libnvvm-sys" } nvjitlink-sys = { path = "crates/nvjitlink-sys" } reserved-oxide-symbols = { path = "crates/reserved-oxide-symbols" } -fuzzer = { path = "crates/fuzzer" } # The crates.io release, not the in-tree path: the shared cuda-core (cutile-rs) # reads artifact bundles through crates.io oxide-artifacts, and a path copy of # the same version is a second crate to rustc. crates/oxide-artifacts @@ -98,19 +113,27 @@ fuzzer = { path = "crates/fuzzer" } # publish, then move this line. # scripts/check-oxide-artifacts-parity.sh keeps them equal. oxide-artifacts = "0.2.1" +# Tile export descriptions shared by `#[tile_export]`, the backend, and +# `cargo oxide tilelib`. +oxide-tile-export = { path = "crates/oxide-tile-export" } +# Tile IR bytecode for those descriptions: `cargo oxide tilelib` writes a +# tile library's interface and implementation halves with it. Pure Rust, so +# cargo-oxide stays toolkit-free. +oxide-tilebc = { path = "crates/oxide-tilebc" } ptx-parse = { path = "crates/ptx-parse" } ptx-schedule = { path = "crates/ptx-schedule" } # Common dependencies -futures = "0.3" thiserror = "2" linkme = "0.3" anyhow = "1.0" once_cell = "1.18" -combine = "4.6" +# Pliron 0.19 requires SourcePosition: Hash, introduced in combine 4.6.8. +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 36d083ba5c..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 @@ -61,10 +65,11 @@ test: -p cuda-intrinsics-gen -p cuda-target-spec -p cuda-intrinsics -p llvm-export \ -p dialect-mir -p dialect-nvvm -p mir-importer -p mir-lower \ -p mir-transforms -p nvvm-transforms -p reserved-oxide-symbols \ - -p cuda-device -p libnvvm-sys -p nvjitlink-sys \ + -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/LICENSE b/cuda-oxide/LICENSE deleted file mode 100644 index f3fe76ecad..0000000000 --- a/cuda-oxide/LICENSE +++ /dev/null @@ -1,203 +0,0 @@ -Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/cuda-oxide/README.md b/cuda-oxide/README.md index c27e17d2fd..3d34f985f3 100644 --- a/cuda-oxide/README.md +++ b/cuda-oxide/README.md @@ -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-08-28 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-08-28 -rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain nightly-2026-08-28 +rustup toolchain install 1.99.1 +rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain 1.99.1 ``` #### CUDA @@ -315,6 +315,9 @@ 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!`) | @@ -337,6 +340,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, ...) | @@ -351,6 +356,8 @@ cargo oxide run gemm_sol_final | `cuda-intrinsics-gen` | Extractor and deterministic source generator for the intrinsics | | `cuda-artifact-finalizer` | Driver-independent NVVM IR, LTOIR, and PTX finalization | | `oxide-artifacts` | Architecture-neutral embedded device artifact metadata | +| `oxide-tile-export` | Export descriptions shared by `cuda-macros`, the backend and `cargo-oxide`; part of the Tile IR interop work, not released yet | +| `oxide-tilebc` | Tile IR bytecode writer for those export descriptions | | `reserved-oxide-symbols` | Workspace-private `cuda_oxide_*` symbol-name contract | | `fuzzer` | Differential codegen fuzzer support (rustlantis adapter) | | `ptx-schedule` | PTX schedule-perturbation fuzzing (nanosleep injection campaigns) | @@ -391,6 +398,6 @@ cuda-oxide is one of several Rust + GPU efforts under active development. Projec ## License -cuda-oxide is licensed under the Apache License, Version 2.0: [LICENSE](LICENSE). +cuda-oxide is licensed under the Apache License, Version 2.0: [LICENSE](../LICENSE). Third-party components retain the licenses stated in their files; see [dependency-licenses.csv](dependency-licenses.csv) for the tracked license inventory. diff --git a/cuda-oxide/THIRD_PARTY_NOTICES b/cuda-oxide/THIRD_PARTY_NOTICES index 82e2ed68c9..df81ba851f 100644 --- a/cuda-oxide/THIRD_PARTY_NOTICES +++ b/cuda-oxide/THIRD_PARTY_NOTICES @@ -4,11 +4,63 @@ THIRD PARTY NOTICES This file attributes third-party open source software whose source code is embedded in cuda-oxide. -cuda-oxide's own code is licensed under Apache-2.0 (see LICENSE). Third-party +cuda-oxide's own code is licensed under Apache-2.0 (see the LICENSE at the +repository root). Third-party 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 -------------------------------------------------------------------------------- @@ -29,7 +81,7 @@ The upstream license texts are retained in place at carries its own identical copies at `crates/fuzzer/rustlantis/rangemap/LICENSE-APACHE` and `crates/fuzzer/rustlantis/rangemap/LICENSE-MIT`. The Apache-2.0 text is -identical to the one in this directory (LICENSE). The MIT text is +identical to the repository's root LICENSE. The MIT text is reproduced here: Permission is hereby granted, free of charge, to any diff --git a/cuda-oxide/crates/cargo-oxide/Cargo.toml b/cuda-oxide/crates/cargo-oxide/Cargo.toml index 7c34766a0c..45e2a16cbe 100644 --- a/cuda-oxide/crates/cargo-oxide/Cargo.toml +++ b/cuda-oxide/crates/cargo-oxide/Cargo.toml @@ -23,6 +23,8 @@ cuda-artifact-finalizer = { workspace = true } libnvvm-sys = { workspace = true } nvjitlink-sys = { workspace = true } oxide-artifacts = { workspace = true } +oxide-tile-export = { workspace = true } +oxide-tilebc = { workspace = true } ptx-parse = { workspace = true } ptx-schedule = { workspace = true } reserved-oxide-symbols = { workspace = true } diff --git a/cuda-oxide/crates/cargo-oxide/README.md b/cuda-oxide/crates/cargo-oxide/README.md index f8e685eead..a381dfc581 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-08-28 install --git https://github.com/NVlabs/cuda-oxide.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. @@ -28,6 +28,7 @@ cargo oxide build vecadd # compile only (no run) cargo oxide build vecadd --debug-assertions # keep release optimization, compile debug assertions cargo oxide build vecadd --materialize-cubin --arch sm_120 # embed native cubin cargo oxide emit-ltoir vecadd --arch sm_100 # device code -> .ltoir (Tile/SIMT interop) +cargo oxide tilelib tilelib_relu --arch sm_120 # #[tile_export] crate -> cuTile tile library cargo oxide build -- -p my_app # arbitrary cargo build through cuda-oxide cargo oxide test # cargo test with cuda-oxide defaults cargo oxide test -- -p my_app # arbitrary cargo test through cuda-oxide @@ -53,20 +54,22 @@ cargo oxide update --force # inside the workspace, run setup via update |------------------------------|----------------------------------|-------------------------------------------------| | `--materialize-cubin` | run, sanitize, build, test, pipeline, debug | Finalize and embed target-specific native GPU code during the host build | | `--emit-nvvm-ir` | run, build, pipeline | Generate NVVM IR for libNVVM | -| `--arch ` | run, sanitize, build, test, pipeline, emit-ltoir, inspect, debug | Target architecture override | -| `--features ` | run, sanitize, debug, build, build passthrough, emit-ltoir, inspect | Comma-separated cargo features to enable | +| `--arch ` | run, sanitize, build, test, pipeline, emit-ltoir, tilelib, inspect, debug | Target architecture override | +| `--features ` | run, sanitize, debug, build, build passthrough, emit-ltoir, tilelib, inspect | Comma-separated cargo features to enable | | `--device-features ` | run, build | Comma-separated features for metadata-declared device crates only | | `--bin ` | run, sanitize, debug | Specific binary target to build/run | | `--tool ` | sanitize | Compute Sanitizer tool: `memcheck`, `racecheck`, `initcheck`, or `synccheck` | -| `-o, --output

` | emit-ltoir | Output path for the `.ltoir` artifact | +| `-o, --output

` | emit-ltoir, tilelib | Output path for the `.ltoir` or `.tilelib` artifact | +| `--library-optional` | tilelib | Succeed when the installed tileiras cannot build the `.tilelib`; the descriptions and LTOIR are still written | +| `--print` | tilelib | Print the readable form of the interface and implementation modules | | `--cargo-target-dir ` | build/test passthrough | Cargo target directory | | `--device-codegen-crate ` | build/test passthrough | Comma-separated device owner crate filter | | `--device-cfg ` | build/test passthrough | Append `--cfg NAME` to rustflags | -| `-v, --verbose` | run, sanitize, build, test, emit-ltoir, inspect | Show detailed compilation output | -| `--no-fmad` | run, sanitize, build, test, emit-ltoir, pipeline, inspect | Keep ordinary multiply and add/subtract operations separate | -| `--unchecked-indexing` | run, sanitize, build, test, emit-ltoir, pipeline, inspect | Elide device slice/array bounds checks (UB on OOB) | -| `--lineinfo` | run, sanitize, build, test, emit-ltoir, pipeline, inspect | Emit device line-number info for profilers (nvcc `-lineinfo`) | -| `--device-debug` | run, sanitize, build, test, emit-ltoir, pipeline, inspect | Emit full device debug info; disables libNVVM optimization (nvcc `-G`) | +| `-v, --verbose` | run, sanitize, build, test, emit-ltoir, tilelib, inspect | Show detailed compilation output | +| `--no-fmad` | run, sanitize, build, test, emit-ltoir, tilelib, pipeline, inspect | Keep ordinary multiply and add/subtract operations separate | +| `--unchecked-indexing` | run, sanitize, build, test, emit-ltoir, tilelib, pipeline, inspect | Elide device slice/array bounds checks (UB on OOB) | +| `--lineinfo` | run, sanitize, build, test, emit-ltoir, tilelib, pipeline, inspect | Emit device line-number info for profilers (nvcc `-lineinfo`) | +| `--device-debug` | run, sanitize, build, test, emit-ltoir, tilelib, pipeline, inspect | Emit full device debug info; disables libNVVM optimization (nvcc `-G`) | | `--debug-assertions` | build | Compile `debug_assert!` / `cfg(debug_assertions)` with release-like optimization and overflow checks disabled | | `--force` | update | Inside the workspace, run `setup` instead of advising it | | `--async` | new | Use the async template | @@ -75,8 +78,8 @@ cargo oxide update --force # inside the workspace, run setup via update | `--check` | fmt | Check formatting only | | `--json` | list | Emit stable machine-readable output | -`--arch` is required for `emit-ltoir` and explicit NVVM IR output because those -artifacts are architecture-specific. Without an override, `run` detects the +`--arch` is required for `emit-ltoir`, `tilelib` and explicit NVVM IR output +because those artifacts are architecture-specific. Without an override, `run` detects the local GPU, while `build` and `pipeline` use the compiler's feature-based target so they remain useful for cross-compilation. @@ -124,15 +127,19 @@ multi-artifact linking is not yet represented by this single-cubin path. Metadata-interop projects keep rejecting this global flag; they can instead request one independently finalized cubin per declared device crate with `artifact-kind = "cubin"`, as described below. Materialization is also a -final-output mode, so it cannot be combined with `--emit-nvvm-ir` or -`emit-ltoir`. +final-output mode, so it cannot be combined with `--emit-nvvm-ir`, +`emit-ltoir` or `tilelib`. ## Commands ### `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 @@ -254,10 +261,25 @@ cargo oxide build tcgen05 # sm_100a only, but PTX generation works anywhe cargo oxide build debug --debug-assertions ``` -`--debug-assertions` keeps the normal release-like `-Copt-level=3` policy but -switches to `-Cdebug-assertions=on`. It also explicitly sets -`-Coverflow-checks=off`: rustc otherwise enables overflow checks with debug -assertions, which adds distinct MIR checks and trap paths. The option applies +Builds pin a release-like codegen policy: `opt-level = 3`, +`debug-assertions = off` and `overflow-checks = off` (overflow checks add +distinct MIR checks and trap paths). The pins travel in the Cargo profile, so +a build whose pins match the release defaults shares every compiled crate +with `cargo oxide test -- --release`. Two carriers deliver them: the `release` +and `dev` profiles take `CARGO_PROFILE_RELEASE_*` and `CARGO_PROFILE_DEV_*` +environment settings, and a profile you select with `--profile` takes +`--config profile..=` arguments placed before your own +Cargo arguments. When `RUSTFLAGS`, `CARGO_ENCODED_RUSTFLAGS` or the project's +`extra-rustflags` set one of these options, cargo-oxide also appends the +pinned `-C` flag after them, so the pin still wins there. Three channels +outrank the pins: `[profile.*.package.*]` overrides, `--config profile.*` +arguments, and manifest `[profile.*.rustflags]` or +`[profile.*.package.*.rustflags]` lists, which cargo-oxide never sees. Host +units (build scripts, proc-macros) follow Cargo's `build-override` settings. + +`--debug-assertions` keeps `opt-level = 3` and `overflow-checks = off` but +switches to `debug-assertions = on` (rustc would otherwise enable overflow +checks with it). The option applies to regular builds, passthrough builds, and metadata-declared interop device crates. It is independent of `--device-debug`, which controls emitted device debug information and CUDA finalization optimization. @@ -314,6 +336,64 @@ Pre-Blackwell targets use LLVM 7 typed pointers; `compute_100` and newer targets use modern opaque pointers. The target is checked before export and recorded alongside the artifact so the runtime uses the same format. +### `cargo oxide tilelib ` + +Packages a crate's `#[tile_export]` functions as a cuTile tile library: the +other direction of the interop, where a cuTile program calls Rust device +functions through `foreign_call`. + +```text +emit-ltoir ──> .ltoir .ll .tileexports.json + │ │ + define @(...) checked against the ABI table + │ + descriptions ──> oxide-tilebc Declaration + Implementation + │ + .iface.tilebc .impl.tilebc .tilelib.mlir + │ + tileiras --preview-create-tile-lib (lists 13.5 or newer) ──> .tilelib +``` + +The command runs `emit-ltoir` for the device half, then checks that the +emitted NVVM IR defines every implementation symbol the descriptions name +with the parameter list the native SIMT ABI predicts (`in` scalars by value, +everything else a pointer, `void` return), because the assembler never looks +inside the code payload and a missing symbol would otherwise surface only +when a cuTile program is compiled against the library. The descriptions then +become the two halves of the library through the `oxide-tilebc` crate: the +interface (`preview$func extern` declarations with tile-typed arguments) and +the implementation catalog (`preview$extern_impl` entries whose fragments +describe each parameter's buffer and its map onto the tile). Both are +written as Tile IR bytecode with the preview tag, at the version the resolved +assembler calls current (its newest listed version, 13.5 when none can read +them yet), next to a readable `.tilelib.mlir`; `--print` shows that +text. + +`tileiras --preview-create-tile-lib` exists from the CUDA 13.5 assembler on, +the first one whose `--list-versions` names bytecode 13.5. An assembler reads +preview content only at its current version, which is why the descriptions +are written at that version. When the resolved tileiras lists 13.5 or newer, +the library is assembled from the interface, the implementation catalog and +the LTOIR into `.tilelib` (or `-o`). +Otherwise the command explains what the installed assembler can read, points +at the files a newer one needs, and exits 2; `--library-optional` makes that +outcome exit 0, for toolkits that can build the descriptions but not the +library yet. The summary ends with one of two lines a script can read instead +of probing tileiras itself: `Tile library written to ` or +`Tile library skipped: lists bytecode versions ...`. The assembler +is resolved from `CUTILE_TILEIRAS_PATH` (cuTile's own override), then like +every other toolkit tool: PATH, the configured CUDA toolkit root, the standard +install locations. A tileiras that is configured but cannot be run is an +error (exit 1) naming the path and the OS error; only an assembler that runs +and lists no bytecode versions is treated as one that predates +`--list-versions`. + +```bash +cargo oxide tilelib tilelib_relu --arch sm_120 --print +cargo oxide tilelib tilelib_relu --arch sm_120 --library-optional # CUDA 13.4: descriptions and LTOIR +cargo oxide tilelib my_lib --arch sm_120 -o build/my_lib.tilelib +``` + ### `cargo oxide test` Runs `cargo test` through the cuda-oxide backend. With no extra arguments it @@ -372,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 @@ -418,23 +498,23 @@ When `cargo oxide` needs the `librustc_codegen_cuda.so` backend, it searches in Outside the repository the backend follows the project's `Cargo.lock`: ```text -Cargo.lock cuda-device = git+https://github.com/NVlabs/cuda-oxide.git#a1b4f118... +Cargo.lock cuda-device = git+https://github.com/NVIDIA/cuda-rust.git# │ ▼ -~/.cargo/git/checkouts/cuda-oxide-*/a1b4f11/ Cargo's checkout of that commit - │ crates/rustc-codegen-cuda → cargo build +~/.cargo/git/checkouts/cuda-rust-*// Cargo's checkout of that commit + │ cuda-oxide/crates/rustc-codegen-cuda → cargo build ▼ -~/.cargo/cuda-oxide/librustc_codegen_cuda.so + source-rev.txt = a1b4f118... +~/.cargo/cuda-oxide/librustc_codegen_cuda.so + source-rev.txt = ``` 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). @@ -505,7 +585,13 @@ Do not put `RUSTFLAGS` or `CARGO_ENCODED_RUSTFLAGS` in the `[env]` table; use inherited user flags, explicit `--device-cfg` values, and its required compiler flags in boundary-preserving `CARGO_ENCODED_RUSTFLAGS`. Required compiler flags are applied last so inherited settings cannot replace the cuda-oxide backend or -disable its correctness-critical codegen options. +disable its correctness-critical codegen options. Pinned profile options +(opt-level, debug assertions, overflow checks, debug info) travel in the Cargo +profile, as `CARGO_PROFILE_RELEASE_*` and `CARGO_PROFILE_DEV_*` settings for +the base profiles and as `--config profile..*` arguments for a selected +`--profile`. They appear here only when `RUSTFLAGS`, `CARGO_ENCODED_RUSTFLAGS` +or `extra-rustflags` set the same option. Manifest `[profile.*.rustflags]` +lists never pass through here and outrank the pins. ## Architecture diff --git a/cuda-oxide/crates/cargo-oxide/src/backend.rs b/cuda-oxide/crates/cargo-oxide/src/backend.rs index 321e0b5aaa..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." @@ -875,7 +875,8 @@ fn invalidate_cache(cache_dir: &Path) { clear_cache_contents(cache_dir); } -/// Drop the cached `.so` and source tree used by external projects. +/// Drop the cached `.so` and source tree used by external projects, and any +/// temporary an interrupted install left beside them. /// /// Used by `cargo oxide update` for an explicit refresh, distinct from the /// automatic upgrade path in [`invalidate_cache`]. @@ -886,9 +887,27 @@ pub fn clear_cached_backend() -> Option { cache_dir.display() ); clear_cache_contents(&cache_dir); + remove_stranded_backend_temporaries(&cache_dir); Some(cache_dir) } +/// Removes the temporaries left by installs killed between their copy and +/// their rename (the names [`temporary_backend_path`] produces); each one is +/// as large as the backend. Only the explicit `cargo oxide update` path +/// removes them. The automatic invalidation in [`invalidate_cache`] runs from +/// any project's lookup, possibly beside another project's installer, whose +/// live temporary must stay until its rename. +fn remove_stranded_backend_temporaries(cache_dir: &Path) { + let Ok(entries) = std::fs::read_dir(cache_dir) else { + return; + }; + for entry in entries.flatten() { + if is_backend_temporary(&entry.file_name().to_string_lossy()) { + let _ = std::fs::remove_file(entry.path()); + } + } +} + fn clear_cache_contents(cache_dir: &Path) { let _ = std::fs::remove_file(cache_dir.join("librustc_codegen_cuda.so")); let _ = std::fs::remove_dir_all(cache_dir.join("src")); @@ -1037,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)); @@ -1247,7 +1267,7 @@ fn auto_fetch_and_build() -> PathBuf { "clone", "--depth", "1", - "https://github.com/NVlabs/cuda-oxide.git", + "https://github.com/NVIDIA/cuda-rust.git", src_dir.to_str().unwrap(), ]) .status() @@ -1277,6 +1297,12 @@ fn auto_fetch_and_build() -> PathBuf { /// Copies a freshly built backend into `cache_dir` and records the toolchain /// fingerprint and the source commit beside it. /// +/// The `.so` is copied to a temporary file and renamed over the old one, so a +/// `rustc` that already mapped the old file keeps reading it. On error the +/// temporary file is removed and the old `.so` stays. A process killed between +/// the copy and the rename leaves the temporary behind; `cargo oxide update` +/// removes it (see [`remove_stranded_backend_temporaries`]). +/// /// The fingerprint must be written whenever the `.so` is. A `.so` installed /// without one falls back to the mtime checks, which cannot see a toolchain /// swap, so the next lookup would load a backend linked against the wrong @@ -1296,16 +1322,42 @@ fn install_backend_into( ) -> std::io::Result { std::fs::create_dir_all(cache_dir)?; let so_path = cache_dir.join("librustc_codegen_cuda.so"); - // Drop the old record before the copy and write the new one after it, so - // the only intermediate state is "no record", which counts as stale: a - // copy interrupted halfway can never be handed out under the old commit. + let temporary = temporary_backend_path(cache_dir); + if let Err(error) = std::fs::copy(built_so, &temporary) { + let _ = std::fs::remove_file(&temporary); + return Err(error); + } + // The copy leaves the old `.so` and its record untouched, so a copy that + // fails costs nothing. Drop the old record just before the rename and + // write the new one after it, so the only intermediate state is "no + // record", which counts as stale: the new `.so` is never handed out under + // the old commit. record_source_rev(cache_dir, None); - std::fs::copy(built_so, &so_path)?; + if let Err(error) = std::fs::rename(&temporary, &so_path) { + let _ = std::fs::remove_file(&temporary); + return Err(error); + } write_toolchain_fingerprint(cache_dir, build_dir); record_source_rev(cache_dir, source_rev); Ok(so_path) } +/// A temporary name beside the `.so` (same filesystem, so the rename is +/// atomic). The pid and a random part keep two installers apart. +fn temporary_backend_path(cache_dir: &Path) -> PathBuf { + use std::hash::{BuildHasher, RandomState}; + let nonce = RandomState::new().hash_one(()); + cache_dir.join(format!( + "librustc_codegen_cuda.so.{}.{nonce:016x}.tmp", + std::process::id() + )) +} + +/// True for the names [`temporary_backend_path`] produces. +fn is_backend_temporary(name: &str) -> bool { + name.starts_with("librustc_codegen_cuda.so.") && name.ends_with(".tmp") +} + /// What [`publish_to_cache`] installed. pub struct PublishedBackend { /// The cached `.so`. @@ -1579,6 +1631,24 @@ mod tests { ); } + /// `clear_cache_contents` doubles as the automatic invalidation that any + /// project's lookup can run, possibly beside another project's installer, + /// so it leaves that installer's temporary alone. Only `cargo oxide + /// update` removes temporaries. + #[test] + fn clear_cache_contents_leaves_a_live_temporary_alone() { + let dir = tempdir(); + let live = temporary_backend_path(&dir); + std::fs::write(&live, b"half-copied").unwrap(); + + clear_cache_contents(&dir); + + assert!( + live.exists(), + "an automatic invalidation must not remove another installer's temporary" + ); + } + /// A backend source input newer than the cached `.so` must report /// `StaleVsSource` (the "developer advanced the source" case that issue /// #49's binary check alone misses). To isolate the source signal from @@ -1867,6 +1937,166 @@ mod tests { ); } + /// The new backend is shorter, so an in-place rewrite would also truncate + /// the file under the open handle. + #[cfg(unix)] + #[test] + fn installing_leaves_a_loaded_backend_intact() { + use std::io::{Read, Seek}; + use std::os::unix::fs::{MetadataExt, PermissionsExt}; + + const OLD: &str = "old backend, longer than the new one"; + const NEW: &str = "new backend"; + let dir = tempdir(); + let cache = dir.join("cache"); + let old = dir.join("old.so"); + let new = dir.join("new.so"); + std::fs::write(&old, OLD).unwrap(); + std::fs::write(&new, NEW).unwrap(); + std::fs::set_permissions(&old, std::fs::Permissions::from_mode(0o644)).unwrap(); + std::fs::set_permissions(&new, std::fs::Permissions::from_mode(0o755)).unwrap(); + + let installed = + install_backend_into(&cache, &old, &dir, None).expect("install must succeed"); + let mut loaded = std::fs::File::open(&installed).unwrap(); + let mut seen = String::new(); + loaded.read_to_string(&mut seen).unwrap(); + assert_eq!(seen, OLD); + + install_backend_into(&cache, &new, &dir, None).expect("install must succeed"); + + seen.clear(); + loaded.rewind().unwrap(); + loaded.read_to_string(&mut seen).unwrap(); + assert_eq!( + seen, OLD, + "a handle opened before the install must keep reading the old backend" + ); + assert_eq!( + std::fs::read_to_string(&installed).unwrap(), + NEW, + "the path must yield the new backend" + ); + let metadata = std::fs::metadata(&installed).unwrap(); + assert_ne!( + metadata.ino(), + loaded.metadata().unwrap().ino(), + "the new backend must be a new file, not the old inode rewritten" + ); + assert_eq!( + metadata.permissions().mode() & 0o777, + 0o755, + "the installed backend must take the built file's mode" + ); + // Closed first: a network filesystem keeps a replaced file under a + // hidden name while a handle to it is open. + drop(loaded); + assert_eq!( + unexpected_cache_entries(&cache), + Vec::::new(), + "installing must not leave a temporary file in the cache" + ); + } + + /// The copy fails before anything in the cache changes, so the backend + /// and the commit recorded for it both stay. + #[test] + fn a_failed_install_keeps_the_cached_backend() { + let dir = tempdir(); + let cache = dir.join("cache"); + let source = dir.join("built.so"); + std::fs::write(&source, b"built").unwrap(); + let installed = install_backend_into(&cache, &source, &dir, Some("aaaa1111")) + .expect("install must succeed"); + + let error = install_backend_into(&cache, &dir.join("missing.so"), &dir, None) + .expect_err("installing a build that does not exist must fail"); + assert_eq!(error.kind(), std::io::ErrorKind::NotFound); + assert_eq!( + std::fs::read(&installed).unwrap(), + b"built", + "a failed install must keep the cached backend" + ); + assert_eq!( + recorded_source_rev(&cache).as_deref(), + Some("aaaa1111"), + "a failed copy must keep the record of the backend that is still there" + ); + assert_eq!( + unexpected_cache_entries(&cache), + Vec::::new(), + "a failed install must not leave a temporary file in the cache" + ); + } + + /// A process killed between the copy and the rename leaves its temporary + /// behind. A later install leaves it alone, since it cannot tell it from + /// another installer's live copy; `cargo oxide update` removes it. + #[test] + fn update_removes_a_temporary_left_by_a_killed_install() { + let dir = tempdir(); + let cache = dir.join("cache"); + let source = dir.join("built.so"); + std::fs::write(&source, b"built").unwrap(); + let installed = install_backend_into(&cache, &source, &dir, Some("aaaa1111")) + .expect("install must succeed"); + // One named by a process that is gone and one named by this process, + // so the sweep and the name policy are checked against each other. + let stranded = [ + cache.join("librustc_codegen_cuda.so.1.00000000deadbeef.tmp"), + temporary_backend_path(&cache), + ]; + for path in &stranded { + std::fs::write(path, b"half-copied").unwrap(); + } + + install_backend_into(&cache, &source, &dir, Some("aaaa1111")) + .expect("install must succeed"); + for path in &stranded { + assert!( + path.exists(), + "an install must not remove what may be another installer's temporary" + ); + } + + remove_stranded_backend_temporaries(&cache); + for path in &stranded { + assert!( + !path.exists(), + "`cargo oxide update` must remove a stranded temporary" + ); + } + assert_eq!( + std::fs::read(&installed).unwrap(), + b"built", + "the sweep must leave the backend alone" + ); + assert_eq!( + recorded_source_rev(&cache).as_deref(), + Some("aaaa1111"), + "the sweep must leave the record alone" + ); + assert_eq!(unexpected_cache_entries(&cache), Vec::::new()); + } + + /// The backend's name is taken by a directory, so the rename fails. + #[test] + fn a_failed_rename_removes_the_temporary_copy() { + let dir = tempdir(); + let cache = dir.join("cache"); + std::fs::create_dir_all(cache.join("librustc_codegen_cuda.so").join("occupied")).unwrap(); + let source = dir.join("built.so"); + std::fs::write(&source, b"built").unwrap(); + + install_backend_into(&cache, &source, &dir, None) + .expect_err("installing over a directory must fail"); + assert_eq!( + unexpected_cache_entries(&cache), + Vec::::new(), + "a failed install must not leave its temporary copy in the cache" + ); + } + /// A cached `.so` whose recorded toolchain fingerprint differs from the /// active toolchain must be `StaleVsToolchain`, even when the mtimes alone /// would call it fresh. This is the case the mtime checks miss: the active @@ -2579,4 +2809,20 @@ mod tests { f.write_all(contents).unwrap(); f.set_modified(mtime).unwrap(); } + + /// Names in `cache` other than the backend and the two records beside it. + fn unexpected_cache_entries(cache: &Path) -> Vec { + let known = [ + "librustc_codegen_cuda.so", + TOOLCHAIN_FINGERPRINT_FILE, + SOURCE_REV_FILE, + ]; + let mut names = std::fs::read_dir(cache) + .unwrap() + .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned()) + .filter(|name| !known.contains(&name.as_str())) + .collect::>(); + names.sort(); + names + } } diff --git a/cuda-oxide/crates/cargo-oxide/src/backend_source.rs b/cuda-oxide/crates/cargo-oxide/src/backend_source.rs index b3a9286951..ce5888cf7c 100644 --- a/cuda-oxide/crates/cargo-oxide/src/backend_source.rs +++ b/cuda-oxide/crates/cargo-oxide/src/backend_source.rs @@ -17,14 +17,14 @@ //! dependency instead of from a separately cloned `main`: //! //! ```text -//! Cargo.lock cuda-device = git+https://github.com/NVlabs/cuda-oxide.git# +//! Cargo.lock cuda-device = git+https://github.com/NVIDIA/cuda-rust.git# //! │ cargo metadata: package.manifest_path //! ▼ -//! ~/.cargo/git/checkouts/cuda-oxide-//crates/cuda-device/Cargo.toml -//! │ walk up to the repository root +//! ~/.cargo/git/checkouts/cuda-rust-//cuda-oxide/crates/cuda-device/Cargo.toml +//! │ walk up to the checkout root that holds the backend crate //! ▼ -//! ~/.cargo/git/checkouts/cuda-oxide-//crates/rustc-codegen-cuda built -//! ~/.cargo/git/checkouts/cuda-oxide-//rust-toolchain.toml nightly it needs +//! ~/.cargo/git/checkouts/cuda-rust-//cuda-oxide/crates/rustc-codegen-cuda built +//! ~/.cargo/git/checkouts/cuda-rust-//cuda-oxide/rust-toolchain.toml nightly it needs //! ``` //! //! The checkout's `rust-toolchain.toml` matters as much as the code: the @@ -316,7 +316,7 @@ mod tests { use serde_json::json; use std::time::SystemTime; - const REPO: &str = "https://github.com/NVlabs/cuda-oxide.git"; + const REPO: &str = "https://github.com/NVIDIA/cuda-rust.git"; const SHA: &str = "a1b4f11882592fae9d022c86a9d8d1a4c9426980"; /// Lays out what Cargo's checkout of the repository looks like: the crate diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/artifacts.rs b/cuda-oxide/crates/cargo-oxide/src/commands/artifacts.rs index 85ddd2eaf1..94e2caaa36 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/artifacts.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/artifacts.rs @@ -112,6 +112,15 @@ pub(super) const GENERATED_ARTIFACT_SUFFIXES: &[&str] = &[ "target", "options", "cubin.target", + // `#[tile_export]` descriptions the backend writes next to the `.ll` + // (`oxide_tile_export::DESCRIPTIONS_SUFFIX`). + "tileexports.json", + // What `cargo oxide tilelib` writes from those descriptions: the two + // Tile IR bytecode halves, their readable dump, and the library itself. + TILELIB_INTERFACE_SUFFIX, + TILELIB_IMPLEMENTATION_SUFFIX, + TILELIB_DUMP_SUFFIX, + TILELIB_LIBRARY_SUFFIX, ]; pub(super) fn generated_artifact_paths(project_dir: &Path, package_name: &str) -> Vec { diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/build_run.rs b/cuda-oxide/crates/cargo-oxide/src/commands/build_run.rs index ac51bc3b01..27e7640d2b 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/build_run.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/build_run.rs @@ -118,7 +118,8 @@ pub fn codegen_run( println!(); } println!("This is the proper cargo workflow:"); - println!(" CARGO_ENCODED_RUSTFLAGS= cargo run"); + println!(" CARGO_ENCODED_RUSTFLAGS= CARGO_PROFILE_*= cargo run"); + println!(" (a custom --profile takes the pins as --config profile..=)"); println!(); touch_main_rs(&example_dir); 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 ccb9fbe768..0f35afd0de 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::ffi::{OsStr, OsString}; use std::path::Path; use std::process::Command; @@ -21,11 +22,11 @@ pub(super) const ENCODED_RUSTFLAGS_SEPARATOR: char = '\u{1f}'; pub(super) const FULL_DEBUG_GET_MUT_OUTLINE_CFG: &str = "cuda_oxide_internal_outline_disjoint_get_mut_v1"; -/// Profile-related rustc flags owned by cuda-oxide. +/// Profile-related codegen settings owned by cuda-oxide. /// /// Backend selection and MIR/symbol invariants are always applied separately. -/// `CargoSelected` deliberately adds no optimization, assertion, or debug-info -/// flags so Cargo's chosen profile remains authoritative. +/// `CargoSelected` deliberately pins no optimization, assertion, or debug-info +/// setting so Cargo's chosen profile remains authoritative. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(super) enum CodegenProfilePolicy { CargoSelected, @@ -42,6 +43,197 @@ impl CodegenProfilePolicy { Self::ReleaseLike } } + + /// The codegen options this policy pins, in rustc flag order. + pub(super) fn pinned_options(self) -> &'static [PinnedOption] { + const OPT_LEVEL: PinnedOption = PinnedOption::new("opt-level", "3", "opt-level", "3"); + const NO_ASSERTIONS: PinnedOption = + PinnedOption::new("debug-assertions", "off", "debug-assertions", "false"); + const ASSERTIONS: PinnedOption = + PinnedOption::new("debug-assertions", "on", "debug-assertions", "true"); + // Overflow-check MIR changes code shape enough to break pattern-sensitive + // device lowerings, so it is off in every release-like policy. Pinned + // explicitly: rustc derives it from debug assertions, and Cargo's dev + // profile (passthrough builds without `--release`) would turn it on. + const NO_OVERFLOW_CHECKS: PinnedOption = + PinnedOption::new("overflow-checks", "off", "overflow-checks", "false"); + const DEBUGINFO: PinnedOption = PinnedOption::new("debuginfo", "2", "debug", "2"); + match self { + Self::CargoSelected => &[], + Self::ReleaseLike => &[OPT_LEVEL, NO_ASSERTIONS, NO_OVERFLOW_CHECKS], + Self::ReleaseLikeWithDebugAssertions => &[OPT_LEVEL, ASSERTIONS, NO_OVERFLOW_CHECKS], + Self::ReleaseLikeWithDebugInfo => { + &[OPT_LEVEL, NO_ASSERTIONS, NO_OVERFLOW_CHECKS, DEBUGINFO] + } + } + } +} + +/// One rustc codegen option a release-like policy requires, with its Cargo +/// profile key. +/// +/// The Cargo profile is the primary carrier: Cargo fingerprints resolved +/// profile values, not where they came from, so a release-like `build` whose +/// pins equal the release defaults shares every artifact with a plain +/// `cargo test --release`. A `-C` flag in RUSTFLAGS would change every crate's +/// cache key and rebuild the whole graph between the two. The rustc flag is +/// emitted only when incoming rustflags set the same option, since rustflags +/// follow profile flags on the rustc command line and would otherwise win. +/// +/// Two carriers deliver the pins. The base profiles `release` and `dev` take +/// them from the environment (`CARGO_PROFILE_RELEASE_*` and +/// `CARGO_PROFILE_DEV_*`, see [`apply_profile_pins`]). Any other profile +/// selected with `--profile` takes them as +/// `--config profile..=` arguments (see +/// [`profile_config_args`]), because the environment key for such a name can +/// collide with a base profile's table: Cargo reads +/// `CARGO_PROFILE_RELEASE_LTO_OPT_LEVEL` as `profile.release.lto` and rejects +/// it. +/// +/// Limits of the profile route: `[profile..package.]` overrides, +/// `--config profile.*` arguments and manifest `[profile..rustflags]` or +/// `[profile..package..rustflags]` lists outrank the pins. The +/// trailing `-C` flag covers RUSTFLAGS, `CARGO_ENCODED_RUSTFLAGS` and the +/// project's `extra-rustflags` only; cargo-oxide never sees manifest +/// rustflags. Host units (build scripts, proc-macros) follow Cargo's +/// `build-override` defaults. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct PinnedOption { + /// rustc `-C` option name. + name: &'static str, + /// Value in rustc's `-C name=value` form. + rustc_value: &'static str, + /// Cargo profile key as a manifest writes it (`opt-level`, `debug`). + cargo_key: &'static str, + /// Value in Cargo's profile form. + cargo_value: &'static str, +} + +impl PinnedOption { + const fn new( + name: &'static str, + rustc_value: &'static str, + cargo_key: &'static str, + cargo_value: &'static str, + ) -> Self { + Self { + name, + rustc_value, + cargo_key, + cargo_value, + } + } + + pub(super) fn rustc_flag(self) -> String { + format!("-C{}={}", self.name, self.rustc_value) + } + + /// `CARGO_PROFILE__`. Cargo derives the environment key from + /// the manifest key by uppercasing it and replacing `-` with `_`. + fn cargo_env_key(self, profile: &str) -> String { + format!( + "CARGO_PROFILE_{}_{}", + profile.to_uppercase(), + self.cargo_key.to_uppercase().replace('-', "_") + ) + } + + /// `profile..=`, the value of one `--config` argument. + fn cargo_config_entry(self, profile: &str) -> String { + format!("profile.{profile}.{}={}", self.cargo_key, self.cargo_value) + } +} + +/// Base Cargo profiles that take the pins from the environment: release-like +/// routes run `--release`, passthrough builds without `--release` run `dev`, +/// and every other profile inherits the keys it does not set from one of the +/// two. +const BASE_CARGO_PROFILES: [&str; 2] = ["release", "dev"]; + +fn is_base_cargo_profile(name: &str) -> bool { + BASE_CARGO_PROFILES.contains(&name) +} + +/// Set the policy's pins on the base Cargo profiles through the environment. +/// +/// Only the base profiles can use this carrier. A custom name becomes part of +/// the key as well, and Cargo reads `CARGO_PROFILE_RELEASE_LTO_OPT_LEVEL` as +/// `profile.release.lto`, an error it reports before compiling anything. The +/// profile `--profile` selects gets its pins from [`profile_config_args`]. +pub(super) fn apply_profile_pins(cmd: &mut Command, profile: CodegenProfilePolicy) { + for option in profile.pinned_options() { + for name in BASE_CARGO_PROFILES { + cmd.env(option.cargo_env_key(name), option.cargo_value); + } + } +} + +/// The profile a Cargo argument list selects with `--profile`, as written. +/// +/// Both `--profile ` and `--profile=` are read; the last one +/// counts. Arguments after `--` belong to the program. `--release` and the +/// default select base profiles, which need no lookup. +pub(super) fn selected_cargo_profile(cargo_args: &[String]) -> Option<&str> { + let mut args = cargo_args + .iter() + .map(String::as_str) + .take_while(|arg| *arg != "--"); + let mut selected = None; + while let Some(arg) = args.next() { + if arg == "--profile" { + selected = args.next(); + } else if let Some(name) = arg.strip_prefix("--profile=") { + selected = Some(name); + } + } + selected +} + +/// `--config profile..=` argument pairs that carry the +/// policy's pins to the profile `--profile` selects. +/// +/// Empty for a base profile, which the environment already covers, for a +/// policy without pins, and for a name outside the TOML bare-key alphabet +/// (ASCII letters, digits, `-` and `_`), which cannot be written as an +/// unquoted config key; such a profile keeps the base pins it inherits. The +/// pairs belong right after the subcommand: a later user `--config` for the +/// same key still wins, and `cargo test` never sees them after its `--`. +pub(super) fn profile_config_args(policy: CodegenProfilePolicy, profile: &str) -> Vec { + let bare_key = !profile.is_empty() + && profile + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_'); + if is_base_cargo_profile(profile) || !bare_key { + return Vec::new(); + } + policy + .pinned_options() + .iter() + .flat_map(|option| ["--config".to_string(), option.cargo_config_entry(profile)]) + .collect() +} + +/// Names of the `-C` options a rustflags list sets, normalized to rustc's +/// dashed spelling, including the `-O` (`opt-level`) and `-g` (`debuginfo`) +/// shorthands. +pub(super) fn codegen_option_names(flags: &[String]) -> Vec { + let mut names = Vec::new(); + let mut tokens = flags.iter().map(String::as_str); + while let Some(token) = tokens.next() { + let setting = match token { + "-O" => Some("opt-level"), + "-g" => Some("debuginfo"), + "-C" | "--codegen" => tokens.next(), + _ => token + .strip_prefix("--codegen=") + .or_else(|| token.strip_prefix("-C")), + }; + if let Some(setting) = setting { + let name = setting.split_once('=').map_or(setting, |(name, _)| name); + names.push(name.replace('_', "-")); + } + } + names } /// Construct boundary-preserving rustc flags for Cargo. @@ -99,26 +291,17 @@ pub(super) fn build_encoded_rustflags_with_existing( } flags.extend(explicit_rustflags.iter().cloned()); strip_wrapper_owned_codegen_cfgs(&mut flags); + // Pins travel in the Cargo profile (`apply_profile_pins`); a rustc flag is + // needed only to beat an incoming flag for the same option. + let overridden = codegen_option_names(&flags); + let pin_flags: Vec = profile + .pinned_options() + .iter() + .filter(|option| overridden.iter().any(|name| name == option.name)) + .map(|option| option.rustc_flag()) + .collect(); flags.push(format!("-Zcodegen-backend={}", backend_so.display())); - if matches!( - profile, - CodegenProfilePolicy::ReleaseLike - | CodegenProfilePolicy::ReleaseLikeWithDebugAssertions - | CodegenProfilePolicy::ReleaseLikeWithDebugInfo - ) { - flags.push("-Copt-level=3".to_string()); - if profile == CodegenProfilePolicy::ReleaseLikeWithDebugAssertions { - // rustc normally enables overflow checks together with debug - // assertions. Keep them independent: overflow-check MIR changes - // code shape enough to break pattern-sensitive device lowerings. - flags.extend([ - "-Cdebug-assertions=on".to_string(), - "-Coverflow-checks=off".to_string(), - ]); - } else { - flags.push("-Cdebug-assertions=off".to_string()); - } - } + flags.extend(pin_flags); flags.extend([ "-Zmir-enable-passes=-JumpThreading".to_string(), // Device codegen is whole-program: `collector` walks the call graph from @@ -136,9 +319,6 @@ pub(super) fn build_encoded_rustflags_with_existing( "-Zalways-encode-mir".to_string(), "-Csymbol-mangling-version=v0".to_string(), ]); - if profile == CodegenProfilePolicy::ReleaseLikeWithDebugInfo { - flags.push("-Cdebuginfo=2".to_string()); - } flags.join(&ENCODED_RUSTFLAGS_SEPARATOR.to_string()) } @@ -429,7 +609,9 @@ 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); Ok(()) } @@ -516,6 +698,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; @@ -527,6 +710,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 4168590ce8..a7e419a2c4 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/context.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/context.rs @@ -79,6 +79,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: @@ -587,26 +601,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 17bfaa94d6..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"); @@ -568,7 +562,8 @@ pub fn doctor(ctx: &Context) { Err(e) => { println!("✗ {}", e); eprintln!(" Required only when kernels call CUDA libdevice math."); - eprintln!(" Ships with the CUDA Toolkit at /lib64/libnvJitLink.so."); + eprintln!(" Ships with the CUDA Toolkit at /lib64/libnvJitLink.so,"); + eprintln!(" or /lib/ in conda-forge and nixpkgs layouts."); ok = false; } } @@ -861,8 +856,8 @@ pub(super) struct ToolkitChoice { /// /// Mirrors BY HAND `resolve_cuda_toolkit`, `find_default_cuda_toolkit` and /// `validate_cuda_toolkit` in the shared `cuda-bindings` build script, which -/// lives in NVlabs/cutile-rs (`cuda-bindings/build.rs`) and is pinned at -/// `cuda-bindings = "0.3.1"` in the root manifest: doctor cannot import it +/// lives at the git root of this repository (`cuda-bindings/build.rs`) and is +/// a path dependency of the cuda-oxide workspace: doctor cannot import it /// because build-script logic is not a library. If that discovery changes, /// mirror it here. /// @@ -1018,7 +1013,7 @@ pub(super) fn format_cuda_version(version: u32) -> String { /// all candidates with that single directory, including the standard layout. /// /// Mirrors BY HAND the selection table in the shared `cuda-bindings` build -/// sources in NVlabs/cutile-rs (`cuda-bindings/toolkit_target.rs`, +/// sources at the git root of this repository (`cuda-bindings/toolkit_target.rs`, /// `resolve_toolkit_include_candidates`): doctor cannot import it because /// build-script sources are not a library. If the selection there changes, /// mirror it here. 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/materialize.rs b/cuda-oxide/crates/cargo-oxide/src/commands/materialize.rs index fb49d68aa6..80dee3ee40 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/materialize.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/materialize.rs @@ -310,7 +310,7 @@ pub fn print_materializer_handshake() { }); let handshake = finalizer.materializer_handshake().unwrap_or_else(|| { eprintln!( - "the loaded libNVVM or nvJitLink library cannot be tied to an exact file; refusing materialization because Cargo could not fingerprint the compiler inputs" + "the loaded libNVVM or nvJitLink library cannot be tied to an exact file; refusing materialization because Cargo could not fingerprint the compiler inputs. The usual cause is a library found only by SONAME: set LIBNVVM_PATH or LIBNVJITLINK_PATH to the library file, or CUDA_TOOLKIT_PATH or CUDA_HOME to a toolkit root that contains both" ); std::process::exit(1); }); diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/mod.rs b/cuda-oxide/crates/cargo-oxide/src/commands/mod.rs index f0a94b187d..10bc374ee6 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/mod.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/mod.rs @@ -28,6 +28,7 @@ mod scaffold; mod setup_update; #[cfg(test)] mod tests; +mod tilelib; mod toolkit; use artifacts::*; @@ -47,4 +48,5 @@ pub use passthrough::*; pub use pipeline_debug::*; pub use scaffold::*; pub use setup_update::*; +pub use tilelib::*; use toolkit::*; diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/passthrough.rs b/cuda-oxide/crates/cargo-oxide/src/commands/passthrough.rs index b8c1a13988..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, @@ -225,11 +300,19 @@ pub(super) fn cargo_passthrough_command_with_env( target_arch, &materialization, ); + let policy = cargo_subcommand.codegen_profile(opts.debug_assertions); let mut cmd = Command::new("cargo"); cmd.arg(cargo_subcommand.as_str()); if let Some(features) = opts.features { cmd.args(["--features", features]); } + // A profile other than `release` or `dev` may set a pinned key itself and + // cannot take the environment carrier, so it receives the pins as + // `--config` arguments. They go before the user's arguments: a later user + // `--config` for the same key still wins. + if let Some(profile) = selected_cargo_profile(cargo_args) { + cmd.args(profile_config_args(policy, profile)); + } cmd.args(cargo_args).current_dir(&ctx.workspace_root); // Project configuration provides defaults. Explicit wrapper flags and @@ -242,13 +325,7 @@ pub(super) fn cargo_passthrough_command_with_env( opts.unchecked_indexing, opts.device_debug, ); - apply_codegen_configuration( - &mut cmd, - ctx, - cargo_subcommand.codegen_profile(opts.debug_assertions), - opts.device_cfgs, - &fingerprint, - )?; + apply_codegen_configuration(&mut cmd, ctx, policy, opts.device_cfgs, &fingerprint)?; if let Some(cargo_target_dir) = opts.cargo_target_dir { cmd.env("CARGO_TARGET_DIR", cargo_target_dir); diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/pipeline_debug.rs b/cuda-oxide/crates/cargo-oxide/src/commands/pipeline_debug.rs index 2fe9ea4325..f84a54f918 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/pipeline_debug.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/pipeline_debug.rs @@ -68,9 +68,11 @@ pub fn codegen_show_pipeline( } } println!(); + println!("Required settings (Cargo profile; also rustflags if yours override them):"); + println!(" opt-level=3 MIR optimization"); + println!(" debug-assertions=off Remove debug checks"); + println!(" overflow-checks=off Keep MIR free of overflow trap paths"); println!("Required flags (applied via CARGO_ENCODED_RUSTFLAGS):"); - println!(" -C opt-level=3 MIR optimization"); - println!(" -C debug-assertions=off Remove debug checks"); println!(" -Z mir-enable-passes=-JumpThreading"); println!(" Prevent barrier duplication"); println!(" -Z always-encode-mir Emit MIR for all reachable device deps"); @@ -227,7 +229,6 @@ pub fn codegen_debug( &[], &fingerprint, ); - cmd.env("CARGO_PROFILE_RELEASE_DEBUG", "2"); apply_output_mode(&mut cmd, false, target_arch, &materialization); apply_device_arch_hint(&mut cmd, target_arch, detected_device_arch.as_deref()); apply_ld_library_path(&mut cmd, ctx); diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/scaffold.rs b/cuda-oxide/crates/cargo-oxide/src/commands/scaffold.rs index 702285da1c..8e338985ad 100644 --- a/cuda-oxide/crates/cargo-oxide/src/commands/scaffold.rs +++ b/cuda-oxide/crates/cargo-oxide/src/commands/scaffold.rs @@ -11,19 +11,17 @@ use super::*; // cargo oxide new -- standalone project scaffolding // ========================================================================= -const GIT_REPO: &str = "https://github.com/NVlabs/cuda-oxide.git"; - -/// crates.io release of the host-side runtime (`cuda-core`, `cuda-async`). -/// In-tree examples and `cuda-host` path-depend on the git-root copies; -/// `cargo oxide new` still uses this version because a standalone project -/// is not inside the monorepo. -pub(super) const SHARED_HOST_CRATES_VERSION: &str = "0.4.0"; +const GIT_REPO: &str = "https://github.com/NVIDIA/cuda-rust.git"; const RUST_TOOLCHAIN_TOML: &str = r#"[toolchain] -channel = "nightly-2026-08-28" +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 @@ -58,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, @@ -104,6 +103,12 @@ cargo oxide run ) } +/// Every CUDA crate comes from `GIT_REPO`. `cuda-host` depends on `cuda-core` +/// and `cuda-async` by path at tip of tree, so in a project that takes +/// `cuda-host` from git those crates resolve from the same checkout. Naming +/// them from crates.io instead adds a second, distinct `cuda_core` whose +/// `CudaContext` and `CudaModule` do not unify with the ones `#[cuda_module]` +/// expands against, and the project fails to compile. fn scaffold_cargo_toml(name: &str, async_mode: bool) -> String { if async_mode { format!( @@ -117,9 +122,10 @@ edition = "2024" [dependencies] cuda-device = {{ git = "{GIT_REPO}" }} cuda-host = {{ git = "{GIT_REPO}", features = ["async"] }} -# Shared host-side runtime, published from NVlabs/cutile-rs. -cuda-core = "{SHARED_HOST_CRATES_VERSION}" -cuda-async = "{SHARED_HOST_CRATES_VERSION}" +# Shared host-side runtime. It has to come from the same repository as +# cuda-host, which depends on it; a crates.io copy is a second cuda-core. +cuda-core = {{ git = "{GIT_REPO}" }} +cuda-async = {{ git = "{GIT_REPO}" }} tokio = {{ version = "1", features = ["rt", "rt-multi-thread", "macros"] }} "# ) @@ -135,8 +141,9 @@ edition = "2024" [dependencies] cuda-device = {{ git = "{GIT_REPO}" }} cuda-host = {{ git = "{GIT_REPO}" }} -# Shared host-side runtime, published from NVlabs/cutile-rs. -cuda-core = "{SHARED_HOST_CRATES_VERSION}" +# Shared host-side runtime. It has to come from the same repository as +# cuda-host, which depends on it; a crates.io copy is a second cuda-core. +cuda-core = {{ git = "{GIT_REPO}" }} "# ) } @@ -297,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), @@ -317,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 bf368ab300..148309b068 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 @@ -34,6 +35,7 @@ const COMMANDS_SOURCE: &str = concat!( include_str!("scaffold.rs"), include_str!("setup_update.rs"), include_str!("tests.rs"), + include_str!("tilelib.rs"), include_str!("toolkit.rs"), ); @@ -68,6 +70,50 @@ 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() +} + +fn command_args(cmd: &Command) -> Vec { + cmd.get_args() + .map(|arg| arg.to_string_lossy().into_owned()) + .collect() +} + /// `cargo_passthrough_command` with an empty ambient /// `CUDA_OXIDE_MATERIALIZE_CUBIN`. /// @@ -85,18 +131,19 @@ fn passthrough_command_for_test( cargo_passthrough_command_with_env(ctx, cargo_subcommand, opts, cargo_args, None) } +/// Run one passthrough command on a scratch workspace and report which units +/// Cargo reused, by target name. A test harness carries a ` (test)` suffix so +/// it never shadows the library it was built from. fn cargo_artifact_freshness( ctx: &Context, + cargo_subcommand: CargoPassthroughSubcommand, opts: &CargoPassthroughOptions<'_>, + cargo_args: &[&str], materializer_provenance: Option<&str>, ) -> BTreeMap { - let mut cmd = passthrough_command_for_test( - ctx, - CargoPassthroughSubcommand::Build, - opts, - &["--message-format=json-render-diagnostics".to_string()], - ) - .unwrap(); + let mut args = strings(&["--message-format=json-render-diagnostics"]); + args.extend(strings(cargo_args)); + let mut cmd = passthrough_command_for_test(ctx, cargo_subcommand, opts, &args).unwrap(); if let Some(provenance) = materializer_provenance { // Exercise a non-canonical spelling accepted by the backend. The // macro must still track exact provenance rather than keying that @@ -117,10 +164,13 @@ fn cargo_artifact_freshness( .filter_map(|line| serde_json::from_str::(line).ok()) .filter(|message| message["reason"] == "compiler-artifact") .filter_map(|message| { - Some(( - message["target"]["name"].as_str()?.to_string(), - message["fresh"].as_bool()?, - )) + let name = message["target"]["name"].as_str()?; + let key = if message["profile"]["test"].as_bool()? { + format!("{name} (test)") + } else { + name.to_string() + }; + Some((key, message["fresh"].as_bool()?)) }) .collect() } @@ -501,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)] @@ -1215,15 +1286,23 @@ fn interop_build_keeps_debug_assertions_and_device_debug_independent() { append_full_debug_rustflags(&mut encoded, &cmd, None); let flags = decoded_rustflags(&encoded); - assert!(flags.contains(&"-Copt-level=3")); - assert!(flags.contains(&if assertions { + let pins: Vec = options + .codegen_profile() + .pinned_options() + .iter() + .map(|pin| pin.rustc_flag()) + .collect(); + let assertions_flag = if assertions { "-Cdebug-assertions=on" } else { "-Cdebug-assertions=off" - })); - if assertions { - assert!(flags.contains(&"-Coverflow-checks=off")); - } + }; + assert_eq!( + pins, + ["-Copt-level=3", assertions_flag, "-Coverflow-checks=off"] + ); + // Without incoming overrides the pins stay in the Cargo profile. + assert!(!flags.iter().any(|flag| pins.iter().any(|pin| pin == flag))); assert_eq!(has_full_debug_cfg(&flags), debug == DeviceDebug::Full); for required in FULL_DEBUG_MIR_RUSTFLAGS { assert_eq!(flags.contains(required), debug == DeviceDebug::Full); @@ -2002,12 +2081,12 @@ fn build_passthrough_retains_release_profile_and_required_flags() { assert_eq!(flags[0], "-Lnative=/nix/store/cuda-cudart/lib"); assert!(flags.contains(&"-Copt-level=0")); assert!(flags.contains(&"-Zcodegen-backend=llvm")); + // Only the overridden option is re-pinned; debug-assertions stays in the profile. assert_eq!( - &flags[flags.len() - 6..], + &flags[flags.len() - 5..], [ "-Zcodegen-backend=/tmp/librustc_codegen_cuda.so", "-Copt-level=3", - "-Cdebug-assertions=off", "-Zmir-enable-passes=-JumpThreading", "-Zalways-encode-mir", "-Csymbol-mangling-version=v0", @@ -2068,8 +2147,9 @@ fn encoded_rustflags_preserve_configured_flag_boundaries_and_spaces() { .any(|pair| pair == ["--cfg", "model=\"alpha beta\""]) ); assert_eq!(&flags[2..4], ["-L", "native=/nix/store/cuda-cudart/lib"]); + // Backend, then the three required -Z/-C invariants (no incoming pin overrides). assert_eq!( - flags[flags.len() - 6], + flags[flags.len() - 4], "-Zcodegen-backend=/tmp/backend path/librustc_codegen_cuda.so" ); } @@ -2148,9 +2228,20 @@ fn debug_profile_retains_release_defaults_and_adds_debuginfo() { ); let flags = decoded_rustflags(&rustflags); - assert!(flags.contains(&"-Copt-level=3")); - assert!(flags.contains(&"-Cdebug-assertions=off")); - assert!(flags.contains(&"-Cdebuginfo=2")); + assert_eq!( + CodegenProfilePolicy::ReleaseLikeWithDebugInfo + .pinned_options() + .iter() + .map(|pin| pin.rustc_flag()) + .collect::>(), + [ + "-Copt-level=3", + "-Cdebug-assertions=off", + "-Coverflow-checks=off", + "-Cdebuginfo=2" + ] + ); + assert!(!flags.iter().any(|flag| flag.starts_with("-Cdebuginfo"))); assert!(flags.contains(&"-Zmir-enable-passes=-JumpThreading")); assert!(flags.contains(&"-Zalways-encode-mir")); assert!(flags.contains(&"-Csymbol-mangling-version=v0")); @@ -2569,16 +2660,325 @@ fn build_passthrough_debug_assertions_select_the_wrapper_profile() { let flags = decoded_rustflags(&encoded); assert_eq!( - &flags[flags.len() - 6..], + &flags[flags.len() - 3..], [ - "-Copt-level=3", - "-Cdebug-assertions=on", - "-Coverflow-checks=off", "-Zmir-enable-passes=-JumpThreading", "-Zalways-encode-mir", "-Csymbol-mangling-version=v0", ] ); + for base in ["RELEASE", "DEV"] { + for (key, value) in [ + ("OPT_LEVEL", "3"), + ("DEBUG_ASSERTIONS", "true"), + ("OVERFLOW_CHECKS", "false"), + ] { + assert_eq!( + command_env(&cmd, &format!("CARGO_PROFILE_{base}_{key}")).as_deref(), + Some(value) + ); + } + } +} + +/// A release-like build and `test -- --release` must hand rustc the same +/// rustflags, or Cargo rebuilds every crate between them. +#[test] +fn build_and_test_passthrough_share_rustflags() { + let ctx = test_context(OxideConfig::default()); + let opts = CargoPassthroughOptions { + verbose: false, + emit_nvvm_ir: false, + arch: None, + 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 rustflags = |subcommand| { + let cmd = passthrough_command_for_test(&ctx, subcommand, &opts, &[]).unwrap(); + command_env(&cmd, "CARGO_ENCODED_RUSTFLAGS").unwrap() + }; + assert_eq!( + rustflags(CargoPassthroughSubcommand::Build), + rustflags(CargoPassthroughSubcommand::Test) + ); + let test = + passthrough_command_for_test(&ctx, CargoPassthroughSubcommand::Test, &opts, &[]).unwrap(); + assert!( + !test + .get_envs() + .any(|(key, _)| key.to_string_lossy().starts_with("CARGO_PROFILE_")) + ); +} + +#[test] +fn codegen_option_names_cover_every_rustc_spelling() { + let flags: Vec = [ + "-Copt-level=0", + "-C", + "debug-assertions=on", + "--codegen", + "overflow-checks=on", + "--codegen=debuginfo=1", + "-O", + "-g", + "-Copt_level=1", + "-Cdebug-assertions", + "-Ccodegen-units=1", + "-Lnative=/lib", + "--cfg", + "feature=\"x\"", + "-C", + ] + .map(String::from) + .into(); + assert_eq!( + codegen_option_names(&flags), + [ + "opt-level", + "debug-assertions", + "overflow-checks", + "debuginfo", + "opt-level", + "debuginfo", + "opt-level", + "debug-assertions", + "codegen-units" + ] + ); +} + +/// `--profile` is read in both forms from the Cargo arguments; the last one +/// counts and arguments after `--` belong to the program. +#[test] +fn selected_cargo_profile_reads_both_forms_before_the_separator() { + assert_eq!( + selected_cargo_profile(&strings(&["--profile", "release-lto", "-p", "app"])), + Some("release-lto") + ); + assert_eq!( + selected_cargo_profile(&strings(&["-p", "app", "--profile=ci"])), + Some("ci") + ); + assert_eq!( + selected_cargo_profile(&strings(&["--profile", "first", "--profile=second"])), + Some("second") + ); + assert_eq!( + selected_cargo_profile(&strings(&["--release", "--", "--profile", "ignored"])), + None + ); + assert_eq!(selected_cargo_profile(&strings(&["--profile"])), None); +} + +/// Pins reach the two base profiles through the environment and nothing +/// else: a selected `release-lto` must not add `CARGO_PROFILE_RELEASE_LTO_*`, +/// which Cargo reads as `profile.release.lto` and rejects. +#[test] +fn profile_pins_stay_on_the_base_profiles() { + let ctx = test_context(OxideConfig::default()); + let opts = CargoPassthroughOptions { + verbose: false, + emit_nvvm_ir: false, + arch: None, + 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 cmd = passthrough_command_for_test( + &ctx, + CargoPassthroughSubcommand::Build, + &opts, + &strings(&["--profile", "release-lto"]), + ) + .unwrap(); + let mut expected = BTreeMap::new(); + for base in ["RELEASE", "DEV"] { + for (key, value) in [ + ("OPT_LEVEL", "3"), + ("DEBUG_ASSERTIONS", "false"), + ("OVERFLOW_CHECKS", "false"), + ] { + expected.insert(format!("CARGO_PROFILE_{base}_{key}"), value.to_string()); + } + } + let profile_env: BTreeMap = cmd + .get_envs() + .filter_map(|(key, value)| { + let key = key.to_string_lossy(); + key.starts_with("CARGO_PROFILE_").then(|| { + let value = value.map(|value| value.to_string_lossy().into_owned()); + (key.into_owned(), value.unwrap_or_default()) + }) + }) + .collect(); + assert_eq!(profile_env, expected); + + let mut debug = Command::new("cargo"); + debug.args(["build", "--release"]); + apply_profile_pins(&mut debug, CodegenProfilePolicy::ReleaseLikeWithDebugInfo); + for base in ["RELEASE", "DEV"] { + assert_eq!( + command_env(&debug, &format!("CARGO_PROFILE_{base}_DEBUG")).as_deref(), + Some("2") + ); + } + + let mut test = Command::new("cargo"); + test.args(["test", "--profile=bench"]); + apply_profile_pins(&mut test, CodegenProfilePolicy::CargoSelected); + assert_eq!(test.get_envs().count(), 0); +} + +/// A profile other than `release` or `dev` takes the pins as `--config` +/// arguments placed before the user's Cargo arguments, in both `--profile` +/// forms. Base profiles and `cargo oxide test` add none. +#[test] +fn custom_profile_pins_travel_as_config_arguments() { + let ctx = test_context(OxideConfig::default()); + let opts = CargoPassthroughOptions { + verbose: false, + emit_nvvm_ir: false, + arch: None, + 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 pins = strings(&[ + "--config", + "profile.release-lto.opt-level=3", + "--config", + "profile.release-lto.debug-assertions=false", + "--config", + "profile.release-lto.overflow-checks=false", + ]); + for cargo_args in [ + strings(&["--profile", "release-lto", "-p", "app"]), + strings(&["-p", "app", "--profile=release-lto"]), + ] { + let cmd = passthrough_command_for_test( + &ctx, + CargoPassthroughSubcommand::Build, + &opts, + &cargo_args, + ) + .unwrap(); + let expected: Vec = strings(&["build"]) + .into_iter() + .chain(pins.iter().cloned()) + .chain(cargo_args.iter().cloned()) + .collect(); + assert_eq!(command_args(&cmd), expected); + } + for cargo_args in [ + strings(&[]), + strings(&["--release"]), + strings(&["--profile", "release"]), + strings(&["--profile=dev"]), + ] { + let cmd = passthrough_command_for_test( + &ctx, + CargoPassthroughSubcommand::Build, + &opts, + &cargo_args, + ) + .unwrap(); + let expected: Vec = strings(&["build"]) + .into_iter() + .chain(cargo_args.iter().cloned()) + .collect(); + assert_eq!(command_args(&cmd), expected); + } + + let test = passthrough_command_for_test( + &ctx, + CargoPassthroughSubcommand::Test, + &opts, + &strings(&["--profile", "release-lto"]), + ) + .unwrap(); + assert_eq!( + command_args(&test), + strings(&["test", "--profile", "release-lto"]) + ); + + let with_assertions = CargoPassthroughOptions { + debug_assertions: true, + ..opts + }; + let cmd = passthrough_command_for_test( + &ctx, + CargoPassthroughSubcommand::Build, + &with_assertions, + &strings(&["--profile=ci"]), + ) + .unwrap(); + assert_eq!( + command_args(&cmd), + strings(&[ + "build", + "--config", + "profile.ci.opt-level=3", + "--config", + "profile.ci.debug-assertions=true", + "--config", + "profile.ci.overflow-checks=false", + "--profile=ci", + ]) + ); +} + +/// `profile_config_args` writes every pinned key as a manifest writes it, +/// adds the debug-info pin for the debug route, and leaves base profiles, +/// pin-free policies and names Cargo would reject alone. +#[test] +fn profile_config_args_use_manifest_keys() { + assert_eq!( + profile_config_args( + CodegenProfilePolicy::ReleaseLikeWithDebugInfo, + "release-lto" + ), + strings(&[ + "--config", + "profile.release-lto.opt-level=3", + "--config", + "profile.release-lto.debug-assertions=false", + "--config", + "profile.release-lto.overflow-checks=false", + "--config", + "profile.release-lto.debug=2", + ]) + ); + for (policy, profile) in [ + (CodegenProfilePolicy::ReleaseLike, "release"), + (CodegenProfilePolicy::ReleaseLike, "dev"), + (CodegenProfilePolicy::CargoSelected, "release-lto"), + (CodegenProfilePolicy::ReleaseLike, "my profile"), + (CodegenProfilePolicy::ReleaseLike, ""), + ] { + assert!( + profile_config_args(policy, profile).is_empty(), + "{policy:?} {profile:?}" + ); + } } #[test] @@ -2775,13 +3175,13 @@ device-owner = { path = "../device-owner" } debug_assertions: false, }; - let cold = cargo_artifact_freshness(&ctx, &base, None); + let cold = cargo_artifact_freshness(&ctx, CargoPassthroughSubcommand::Build, &base, &[], None); assert_eq!(cold.get("shared_dep"), Some(&false)); assert_eq!(cold.get("tracked_macro"), Some(&false)); assert_eq!(cold.get("device_owner"), Some(&false)); assert_eq!(cold.get("device-consumer"), Some(&false)); - let warm = cargo_artifact_freshness(&ctx, &base, None); + let warm = cargo_artifact_freshness(&ctx, CargoPassthroughSubcommand::Build, &base, &[], None); assert_eq!(warm.get("shared_dep"), Some(&true)); assert_eq!(warm.get("tracked_macro"), Some(&true)); assert_eq!(warm.get("device_owner"), Some(&true)); @@ -2791,7 +3191,13 @@ device-owner = { path = "../device-owner" } arch: Some("sm_90"), ..base }; - let arch_switch = cargo_artifact_freshness(&ctx, &different_arch, None); + let arch_switch = cargo_artifact_freshness( + &ctx, + CargoPassthroughSubcommand::Build, + &different_arch, + &[], + None, + ); assert_eq!(arch_switch.get("shared_dep"), Some(&true)); assert_eq!(arch_switch.get("tracked_macro"), Some(&true)); assert_eq!(arch_switch.get("device_owner"), Some(&false)); @@ -2801,13 +3207,25 @@ device-owner = { path = "../device-owner" } emit_nvvm_ir: true, ..different_arch }; - let output_switch = cargo_artifact_freshness(&ctx, &different_output, None); + let output_switch = cargo_artifact_freshness( + &ctx, + CargoPassthroughSubcommand::Build, + &different_output, + &[], + None, + ); assert_eq!(output_switch.get("shared_dep"), Some(&true)); assert_eq!(output_switch.get("tracked_macro"), Some(&true)); assert_eq!(output_switch.get("device_owner"), Some(&false)); assert_eq!(output_switch.get("device-consumer"), Some(&false)); - let repeated_output = cargo_artifact_freshness(&ctx, &different_output, None); + let repeated_output = cargo_artifact_freshness( + &ctx, + CargoPassthroughSubcommand::Build, + &different_output, + &[], + None, + ); assert_eq!(repeated_output.get("shared_dep"), Some(&true)); assert_eq!(repeated_output.get("tracked_macro"), Some(&true)); assert_eq!(repeated_output.get("device_owner"), Some(&true)); @@ -2815,7 +3233,9 @@ device-owner = { path = "../device-owner" } let provenance_switch = cargo_artifact_freshness( &ctx, + CargoPassthroughSubcommand::Build, &different_output, + &[], Some("11d91fbe164094f6242d44103d0fb01968b96c6d8f48f124eac8fa73a307a657"), ); assert_eq!(provenance_switch.get("shared_dep"), Some(&true)); @@ -2825,7 +3245,9 @@ device-owner = { path = "../device-owner" } let changed_provenance = cargo_artifact_freshness( &ctx, + CargoPassthroughSubcommand::Build, &different_output, + &[], Some("5b11618c2e44027877d0cd4d0cfd10afed5ef262876791e483ec58f4c5569139"), ); assert_eq!(changed_provenance.get("shared_dep"), Some(&true)); @@ -2836,24 +3258,221 @@ device-owner = { path = "../device-owner" } std::fs::remove_dir_all(root).unwrap(); } -#[test] -fn owner_filter_resolution_is_normalized_and_has_explicit_precedence() { - assert_eq!( - resolve_device_codegen_crates(None, None, Some("gpu-kernels, math_gpu")) - .unwrap() - .as_deref(), - Some("gpu_kernels,math_gpu"), - ); - assert_eq!( - resolve_device_codegen_crates(None, Some("parent-owner"), Some("config-owner")) - .unwrap() - .as_deref(), - Some("parent_owner"), - ); - assert!( - resolve_device_codegen_crates(Some(""), Some("parent-owner"), Some("config-owner")) - .is_err() - ); +/// A workspace of one library and one binary that depends on it, for probing +/// real Cargo. `profiles` is appended to the workspace manifest. The binary +/// computes `u8::MAX + 1`, which panics under overflow checks, so running it +/// shows whether the `overflow-checks = false` pin reached the final crate. +fn write_shared_dependency_workspace(root: &Path, profiles: &str) { + for path in [root.join("shared-dep/src"), root.join("app/src")] { + fs::create_dir_all(path).unwrap(); + } + fs::write( + root.join("Cargo.toml"), + format!( + r#"[workspace] +resolver = "3" +members = ["shared-dep", "app"] +{profiles}"# + ), + ) + .unwrap(); + fs::write( + root.join("shared-dep/Cargo.toml"), + r#"[package] +name = "shared-dep" +version = "0.0.0" +edition = "2024" +"#, + ) + .unwrap(); + fs::write( + root.join("shared-dep/src/lib.rs"), + "pub fn shared_value() -> u32 { 42 }\n", + ) + .unwrap(); + fs::write( + root.join("app/Cargo.toml"), + r#"[package] +name = "app" +version = "0.0.0" +edition = "2024" + +[dependencies] +shared-dep = { path = "../shared-dep" } +"#, + ) + .unwrap(); + fs::write( + root.join("app/src/main.rs"), + r#"fn main() { + assert_eq!(shared_dep::shared_value(), 42); + let value: u8 = std::hint::black_box(u8::MAX); + assert_eq!(value + 1, 0, "overflow checks reached the binary"); +} +"#, + ) + .unwrap(); +} + +fn scratch_workspace_context(root: &Path) -> Context { + Context { + workspace_root: root.to_path_buf(), + codegen_crate: root.join("unused-codegen-source"), + examples_dir: root.join("unused-examples"), + backend_so: PathBuf::from("llvm"), + is_workspace: false, + config: OxideConfig::default(), + } +} + +/// Pins equal to the release defaults keep `build -- --release` and +/// `test -- --release` on one artifact set: the shared dependency is reused +/// when the two alternate. Carried in rustflags, the same pins changed every +/// unit's cache key, so the switch recompiled the whole graph and left a +/// second copy of each artifact behind. +#[test] +fn release_build_and_test_reuse_each_others_dependencies() { + let root = unique_temp_dir("cargo_oxide_release_build_test_reuse"); + let target = root.join("target"); + write_shared_dependency_workspace(&root, ""); + let ctx = scratch_workspace_context(&root); + let opts = CargoPassthroughOptions { + verbose: false, + emit_nvvm_ir: false, + arch: Some("sm_80"), + features: None, + cargo_target_dir: Some(&target), + device_codegen_crate: None, + device_cfgs: &[], + no_fmad: false, + unchecked_indexing: false, + materialize_cubin: false, + device_debug: DeviceDebug::Off, + debug_assertions: false, + }; + + let build = cargo_artifact_freshness( + &ctx, + CargoPassthroughSubcommand::Build, + &opts, + &["--release"], + None, + ); + assert_eq!(build.get("shared_dep"), Some(&false)); + assert_eq!(build.get("app"), Some(&false)); + + let test = cargo_artifact_freshness( + &ctx, + CargoPassthroughSubcommand::Test, + &opts, + &["--release", "--no-run"], + None, + ); + assert_eq!( + test.get("shared_dep"), + Some(&true), + "test -- --release must reuse the release build's dependency" + ); + assert_eq!(test.get("app (test)"), Some(&false)); + + let rebuild = cargo_artifact_freshness( + &ctx, + CargoPassthroughSubcommand::Build, + &opts, + &["--release"], + None, + ); + assert_eq!(rebuild.get("shared_dep"), Some(&true)); + assert_eq!(rebuild.get("app"), Some(&true)); + + fs::remove_dir_all(root).unwrap(); +} + +/// A custom profile with a dash in its name builds in both `--profile` forms, +/// and the pins beat the keys the profile sets itself: `release-lto` asks for +/// overflow checks and `opt-level = 1`, yet the binary wraps instead of +/// panicking. Carried in the environment, the same pins became +/// `CARGO_PROFILE_RELEASE_LTO_*`, which Cargo reads as `profile.release.lto` +/// and rejects before compiling. +#[test] +fn dashed_custom_profile_builds_with_the_pins_in_force() { + let root = unique_temp_dir("cargo_oxide_custom_profile_pins"); + let target = root.join("target"); + write_shared_dependency_workspace( + &root, + r#" +[profile.release-lto] +inherits = "release" +opt-level = 1 +overflow-checks = true +"#, + ); + let ctx = scratch_workspace_context(&root); + let opts = CargoPassthroughOptions { + verbose: false, + emit_nvvm_ir: false, + arch: Some("sm_80"), + features: None, + cargo_target_dir: Some(&target), + device_codegen_crate: None, + device_cfgs: &[], + no_fmad: false, + unchecked_indexing: false, + materialize_cubin: false, + device_debug: DeviceDebug::Off, + debug_assertions: false, + }; + + let built = cargo_artifact_freshness( + &ctx, + CargoPassthroughSubcommand::Build, + &opts, + &["--profile", "release-lto"], + None, + ); + assert_eq!(built.get("app"), Some(&false)); + let status = Command::new(target.join("release-lto").join("app")) + .status() + .expect("failed to run the custom-profile binary"); + assert!( + status.success(), + "overflow checks reached the custom-profile binary: {status}" + ); + + let rebuilt = cargo_artifact_freshness( + &ctx, + CargoPassthroughSubcommand::Build, + &opts, + &["--profile=release-lto"], + None, + ); + assert_eq!( + rebuilt.get("app"), + Some(&true), + "both --profile forms must resolve to one configuration" + ); + + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn owner_filter_resolution_is_normalized_and_has_explicit_precedence() { + assert_eq!( + resolve_device_codegen_crates(None, None, Some("gpu-kernels, math_gpu")) + .unwrap() + .as_deref(), + Some("gpu_kernels,math_gpu"), + ); + assert_eq!( + resolve_device_codegen_crates(None, Some("parent-owner"), Some("config-owner")) + .unwrap() + .as_deref(), + Some("parent_owner"), + ); + assert!( + resolve_device_codegen_crates(Some(""), Some("parent-owner"), Some("config-owner")) + .is_err() + ); } #[test] @@ -2993,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()); @@ -4163,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"); @@ -4768,15 +5592,45 @@ fn scaffold_sync_template_uses_launch_contract_and_docs() { assert!(files.main_rs.contains("prepare_vecadd")); assert!(files.main_rs.contains("LaunchConfig1D")); assert!(!files.main_rs.contains("LaunchConfig::for_num_elems")); - // The host runtime is the crates.io release of the git-root crates. - // In-tree examples path-depend on those crates; a freshly scaffolded - // project is not inside this monorepo, so it cannot use those paths. - assert!( - files + // cuda-core comes from the same git source as cuda-host, which depends on + // it at tip of tree; a crates.io copy would be a second cuda_core. + assert!(files.cargo_toml.contains("cuda-core = { git = \"")); + assert!(!files.cargo_toml.contains("cuda-core = \"")); +} + +/// Every CUDA crate in either template resolves from the git source +/// `cuda-device` names. A host crate taken from crates.io next to a git +/// `cuda-host` puts two `cuda_core` crates in the graph, and `#[cuda_module]` +/// then fails to type-check. +#[test] +fn scaffold_takes_every_cuda_crate_from_one_git_source() { + for async_mode in [false, true] { + let files = scaffold_files("one_source", async_mode); + let git_source = |line: &str| { + line.split("git = \"") + .nth(1) + .and_then(|rest| rest.split('"').next()) + .map(str::to_owned) + }; + let cuda_lines: Vec<&str> = files .cargo_toml - .contains(&format!("cuda-core = \"{SHARED_HOST_CRATES_VERSION}\"")) - ); - assert!(!files.cargo_toml.contains("cuda-core = {")); + .lines() + .filter(|line| line.starts_with("cuda-")) + .collect(); + assert!(cuda_lines.len() >= 3, "{cuda_lines:?}"); + let device = cuda_lines + .iter() + .find(|line| line.starts_with("cuda-device")) + .and_then(|line| git_source(line)) + .expect("cuda-device comes from git"); + for line in cuda_lines { + assert_eq!( + git_source(line).as_deref(), + Some(device.as_str()), + "async={async_mode}: {line}" + ); + } + } } #[test] @@ -4794,13 +5648,10 @@ 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 crates.io; the SIMT surface lives under simt::. - assert!( - files - .cargo_toml - .contains(&format!("cuda-async = \"{SHARED_HOST_CRATES_VERSION}\"")) - ); - assert!(!files.cargo_toml.contains("cuda-async = {")); + // Shared host runtime from the same git source as cuda-host; the SIMT + // surface lives under simt::. + 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!( @@ -4878,6 +5729,23 @@ fn scaffold_gitignore_covers_every_clean_artifact_suffix() { } } +#[test] +fn scaffold_names_the_canonical_repository() { + // Both templates must name the repository the crates live in now; the + // pre-move URL resolves only through GitHub's redirect. + for async_mode in [false, true] { + let files = scaffold_files("demo_kernel", async_mode); + assert!( + files + .cargo_toml + .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\"")); + } +} + #[test] fn device_debug_env_value_matches_the_backend_parser() { // The exported strings must round-trip through the shared @@ -5148,6 +6016,667 @@ fn doctor_backend_source_check_covers_every_branch() { } } +// ============================================================================= +// cargo oxide tilelib +// ============================================================================= + +/// The three exports of the `tilelib_relu` example, as the backend's +/// descriptions file records them. +fn tilelib_relu_descriptions() -> Vec { + use oxide_tile_export::{TileAccess, TileElem, TileExportDesc, TileParamDesc, TileSync}; + let param = |name: &str, access: TileAccess, by_value: bool| TileParamDesc { + name: name.to_string(), + elem: TileElem::F32, + access, + footprint: 1, + by_value, + }; + vec![ + TileExportDesc { + version: oxide_tile_export::CURRENT_VERSION, + decl: "relu".into(), + symbol: "relu_impl".into(), + shape: vec![128], + params: vec![ + param("x", TileAccess::In, true), + param("relu_out", TileAccess::Out, false), + ], + sync: TileSync::None, + arch: Vec::new(), + }, + TileExportDesc { + version: oxide_tile_export::CURRENT_VERSION, + decl: "scale_in_place".into(), + symbol: "scale_in_place_impl".into(), + shape: vec![128], + params: vec![ + param("x", TileAccess::InOut, false), + param("factor", TileAccess::In, true), + ], + sync: TileSync::None, + arch: Vec::new(), + }, + TileExportDesc { + version: oxide_tile_export::CURRENT_VERSION, + decl: "warp_sum".into(), + symbol: "warp_sum_impl".into(), + shape: vec![128], + params: vec![ + param("x", TileAccess::In, true), + param("warp_sum_out", TileAccess::Out, false), + ], + sync: TileSync::Warp, + arch: Vec::new(), + }, + ] +} + +/// Whitespace-insensitive form of a Tile IR text, the rule the upstream +/// fixture tests apply: a whitespace run survives, as one space, only between +/// two word characters, so the hand-written multi-line layout below compares +/// equal to the printer's one-line form while a missing space between two +/// words still differs. +fn normalize_tile_ir(text: &str) -> String { + let is_word = |c: char| c.is_alphanumeric() || "_$.@%!#\"?".contains(c); + let mut out = String::new(); + let mut pending_space = false; + for c in text.chars() { + if c.is_whitespace() { + pending_space = true; + continue; + } + if pending_space { + if let Some(prev) = out.chars().last() + && is_word(prev) + && is_word(c) + { + out.push(' '); + } + pending_space = false; + } + out.push(c); + } + out +} + +/// The interface of `tilelib_relu` in the upstream textual syntax: one +/// extern declaration per export, every argument carrying its access, the +/// results being the out/inout tiles. +const TILELIB_RELU_INTERFACE: &str = r#" +cuda_tile.module @tilelib_relu { + cuda_tile.preview$func extern @relu( + %x: !cuda_tile.tile<128xf32> + {cuda_tile.preview$access = #cuda_tile.preview$access}, + %relu_out: !cuda_tile.tile<128xf32> + {cuda_tile.preview$access = #cuda_tile.preview$access}) + -> !cuda_tile.tile<128xf32> + cuda_tile.preview$func extern @scale_in_place( + %x: !cuda_tile.tile<128xf32> + {cuda_tile.preview$access = #cuda_tile.preview$access}, + %factor: !cuda_tile.tile<128xf32> + {cuda_tile.preview$access = #cuda_tile.preview$access}) + -> !cuda_tile.tile<128xf32> + cuda_tile.preview$func extern @warp_sum( + %x: !cuda_tile.tile<128xf32> + {cuda_tile.preview$access = #cuda_tile.preview$access}, + %warp_sum_out: !cuda_tile.tile<128xf32> + {cuda_tile.preview$access = #cuda_tile.preview$access}) + -> !cuda_tile.tile<128xf32> +} +"#; + +/// The implementation catalog of `tilelib_relu`: by-value scalars are the +/// rank-0 `():()` buffer, pointers the contiguous `(1):(1)` one; the two +/// elementwise exports map one invocation per element and `warp_sum` one +/// element per lane with a whole warp per invocation, which is what its +/// `synchronization = warp` clause requires. +const TILELIB_RELU_IMPLEMENTATION: &str = r#" +cuda_tile.module @tilelib_relu { + cuda_tile.preview$extern_impl @relu_impl implements @relu + (%x: !cuda_tile.preview$fragment<1xf32, local, + mem_layout = #cuda_tile.preview$layout<():()>, + fragmap = #cuda_tile.preview$layout<(128, 1, 1):(1, 0, 0)>> + {cuda_tile.preview$access = #cuda_tile.preview$access}, + %relu_out: !cuda_tile.preview$fragment<1xf32, local, + mem_layout = #cuda_tile.preview$layout<(1):(1)>, + fragmap = #cuda_tile.preview$layout<(128, 1, 1):(1, 0, 0)>> + {cuda_tile.preview$access = #cuda_tile.preview$access}) + cuda_tile.preview$extern_impl @scale_in_place_impl implements @scale_in_place + (%x: !cuda_tile.preview$fragment<1xf32, local, + mem_layout = #cuda_tile.preview$layout<(1):(1)>, + fragmap = #cuda_tile.preview$layout<(128, 1, 1):(1, 0, 0)>> + {cuda_tile.preview$access = #cuda_tile.preview$access}, + %factor: !cuda_tile.preview$fragment<1xf32, local, + mem_layout = #cuda_tile.preview$layout<():()>, + fragmap = #cuda_tile.preview$layout<(128, 1, 1):(1, 0, 0)>> + {cuda_tile.preview$access = #cuda_tile.preview$access}) + cuda_tile.preview$extern_impl @warp_sum_impl implements @warp_sum + (%x: !cuda_tile.preview$fragment<1xf32, local, + mem_layout = #cuda_tile.preview$layout<():()>, + fragmap = #cuda_tile.preview$layout<(4, 32, 1):(32, 1, 0)>> + {cuda_tile.preview$access = #cuda_tile.preview$access}, + %warp_sum_out: !cuda_tile.preview$fragment<1xf32, local, + mem_layout = #cuda_tile.preview$layout<(1):(1)>, + fragmap = #cuda_tile.preview$layout<(4, 32, 1):(32, 1, 0)>> + {cuda_tile.preview$access = #cuda_tile.preview$access}) + synchronization = #cuda_tile.preview$synchronization +} +"#; + +#[test] +fn tilelib_relu_descriptions_print_as_the_upstream_modules() { + let library = TileLibrary::from_descriptions(&tilelib_relu_descriptions()) + .unwrap() + .library("tilelib_relu"); + library.verify().unwrap(); + let interface = library.interface_text(); + assert_eq!( + normalize_tile_ir(&interface), + normalize_tile_ir(TILELIB_RELU_INTERFACE), + "printed interface:\n{interface}" + ); + let implementation = library.implementation_text(); + assert_eq!( + normalize_tile_ir(&implementation), + normalize_tile_ir(TILELIB_RELU_IMPLEMENTATION), + "printed implementation:\n{implementation}" + ); + // Both halves are preview content, written at the version the + // assembler calls current with the header tag set. + let version = oxide_tilebc::BytecodeVersion::PREVIEW_FIRST; + for bytes in [ + library.interface_bytecode(version).unwrap(), + library.implementation_bytecode(version).unwrap(), + ] { + assert_eq!(&bytes[8..12], &[13, 5, 1, 0]); + } + // The readable dump names the file each module ends up in. + let dump = library_dump(&library, "tilelib_relu"); + assert!(dump.starts_with("// tilelib_relu.iface.tilebc"), "{dump}"); + assert!(dump.contains("\n// tilelib_relu.impl.tilebc"), "{dump}"); + assert_eq!(dump.matches("cuda_tile.module @tilelib_relu {").count(), 2); +} + +#[test] +fn fragmaps_follow_the_synchronization_and_the_footprint() { + use oxide_tile_export::TileSync; + use oxide_tilebc::Layout; + // One invocation per element; a footprint of F gives N/F invocations of + // F contiguous elements. + assert_eq!( + fragmap(TileSync::None, 128, 1).unwrap(), + Layout::flat(&[128, 1, 1], &[1, 0, 0]) + ); + assert_eq!( + fragmap(TileSync::None, 128, 4).unwrap(), + Layout::flat(&[32, 1, 4], &[4, 0, 1]) + ); + // A warp per invocation: 32 participants, one element each, or F + // contiguous elements each with adjacent lanes adjacent in the tile. + assert_eq!( + fragmap(TileSync::Warp, 128, 1).unwrap(), + Layout::flat(&[4, 32, 1], &[32, 1, 0]) + ); + assert_eq!( + fragmap(TileSync::Warp, 256, 4).unwrap(), + Layout::flat(&[2, 32, 4], &[128, 4, 1]) + ); + // The footprint must divide the tile size. + assert!( + fragmap(TileSync::None, 128, 3) + .unwrap_err() + .contains("does not divide") + ); + assert!( + fragmap(TileSync::Warp, 64, 4) + .unwrap_err() + .contains("multiple of 128") + ); + assert!( + fragmap(TileSync::None, 128, 0) + .unwrap_err() + .contains("positive") + ); + + // Footprints above one verify against the declaration's tile like the + // one-element case does, when every parameter shares them: `fn + // scale4(x: &mut [f32; 4], factor: &[f32; 4])`. + let mut descriptions = tilelib_relu_descriptions(); + let scale = &mut descriptions[1]; + for param in &mut scale.params { + param.footprint = 4; + param.by_value = false; + } + let library = TileLibrary::from_descriptions(&descriptions) + .unwrap() + .library("lib"); + library.verify().unwrap(); + let text = library.implementation_text(); + assert!( + text.contains( + "%x: !cuda_tile.preview$fragment<4xf32, local, mem_layout = \ + #cuda_tile.preview$layout<(4):(1)>, fragmap = \ + #cuda_tile.preview$layout<(32, 1, 4):(4, 0, 1)>>" + ), + "{text}" + ); + assert!( + text.contains( + "%factor: !cuda_tile.preview$fragment<4xf32, local, mem_layout = \ + #cuda_tile.preview$layout<(4):(1)>, fragmap = \ + #cuda_tile.preview$layout<(32, 1, 4):(4, 0, 1)>>" + ), + "{text}" + ); + // A window beside a scalar cannot share one invocation count, so the + // description is refused before the verifier would reject the fragmaps. + descriptions[1].params[1].footprint = 1; + descriptions[1].params[1].by_value = true; + let error = TileLibrary::from_descriptions(&descriptions).unwrap_err(); + assert!( + error.contains("`x` (footprint 4) and `factor` (footprint 1) differ"), + "{error}" + ); +} + +#[test] +fn a_bf16_scalar_passed_by_value_adds_the_frontend_abi_entry() { + use oxide_tile_export::TileElem; + // cuda-device's bf16 is a u16 wrapper, so the LLVM parameter is i16 + // while the tile element is bf16: the implementation has to say so. + let mut descriptions = tilelib_relu_descriptions(); + let relu = &mut descriptions[0]; + for param in &mut relu.params { + param.elem = TileElem::BF16; + } + let (_, implementation) = describe_export(relu).unwrap(); + assert_eq!( + implementation.frontend_abi, + vec![("bf16".to_string(), "i16".to_string())] + ); + // No `sync` means no clause at all, not an explicit `synchronization`. + assert_eq!(implementation.synchronization, None); + let library = TileLibrary::from_descriptions(&descriptions[..1]) + .unwrap() + .library("lib"); + library.verify().unwrap(); + let text = library.implementation_text(); + assert!( + normalize_tile_ir(&text).contains("{frontendABI={bf16=\"i16\"}}"), + "{text}" + ); + assert!(text.contains("preview$fragment<1xbf16, local"), "{text}"); + + // A bf16 behind a pointer travels as a pointer and needs no entry. + let scale = &mut descriptions[1]; + scale.params[0].elem = TileElem::BF16; + let (_, implementation) = describe_export(scale).unwrap(); + assert!(implementation.frontend_abi.is_empty()); + + // Unsigned Rust integers are the signless Tile IR integers. + let warp_sum = &mut descriptions[2]; + for param in &mut warp_sum.params { + param.elem = TileElem::U32; + } + let (func, implementation) = describe_export(warp_sum).unwrap(); + assert_eq!(func.args[0].ty.element, oxide_tilebc::ScalarType::I32); + assert_eq!( + implementation.args[1].fragment.element, + oxide_tilebc::ScalarType::I32 + ); + assert_eq!( + implementation.synchronization, + Some(oxide_tilebc::Synchronization::Warp) + ); +} + +/// The NVVM IR the backend emits for `tilelib_relu` on a modern target, +/// reduced to its function definitions. +const TILELIB_RELU_LL: &str = "\ +; ModuleID = 'builtin.module' +source_filename = \"tilelib_relu\" +target triple = \"nvptx64-nvidia-cuda\" + +define void @relu_impl(float %v0, ptr %v1) #0 { +entry: + ret void +} + +define void @scale_in_place_impl(ptr %v0, float %v1) #0 { +entry: + ret void +} + +define void @warp_sum_impl(float %v0, ptr %v1) #0 { +entry: + ret void +} + +declare float @llvm.nvvm.shfl.sync.bfly.f32(i32, float, i32, i32) #0 + +define float @tilelib_relu__warp_sum(float %v0) #0 { +entry: + ret float %v0 +} +"; + +#[test] +fn ll_defines_are_read_with_their_return_and_parameter_types() { + let defines = ll_defines(TILELIB_RELU_LL); + assert_eq!( + defines.keys().collect::>(), + [ + "relu_impl", + "scale_in_place_impl", + "tilelib_relu__warp_sum", + "warp_sum_impl" + ] + ); + assert_eq!( + defines["relu_impl"], + LlDefine { + ret: "void".into(), + params: vec!["float".into(), "ptr".into()], + local_linkage: None, + } + ); + assert_eq!(defines["tilelib_relu__warp_sum"].ret, "float"); + // `declare` lines are not definitions. + assert!(!defines.contains_key("llvm.nvvm.shfl.sync.bfly.f32")); + + // Linkage, calling convention and parameter attributes do not change + // the types; address spaces and legacy typed pointers survive intact. A + // linkage that keeps the symbol out of the module's exports is kept + // aside for the symbol check to report. + let defines = ll_defines( + "define internal fastcc void @k(ptr noundef %a, i64 %n, ptr addrspace(3) %s, i8* %p, { ptr, i64 } %slice) #0 {\n\ + define dso_local weak_odr void @w() {\n\ + define private void @p() {\n\ + define linkonce_odr void @l() {\n\ + define ptx_kernel void @\"quoted name\"() {\n", + ); + assert_eq!( + defines["k"], + LlDefine { + ret: "void".into(), + params: vec![ + "ptr".into(), + "i64".into(), + "ptr addrspace(3)".into(), + "i8*".into(), + "{ ptr, i64 }".into(), + ], + local_linkage: Some("internal".into()), + } + ); + assert_eq!(defines["w"].ret, "void"); + assert_eq!(defines["w"].local_linkage, None); + assert_eq!(defines["p"].local_linkage.as_deref(), Some("private")); + assert_eq!(defines["l"].local_linkage.as_deref(), Some("linkonce_odr")); + assert_eq!(defines["quoted name"].params, Vec::::new()); + assert_eq!(defines["quoted name"].local_linkage, None); +} + +#[test] +fn implementation_symbols_are_checked_against_the_abi_table() { + let descriptions = tilelib_relu_descriptions(); + check_implementation_symbols(TILELIB_RELU_LL, &descriptions).unwrap(); + + // The legacy LLVM 7 dialect spells every pointer `i8*`. + let legacy = TILELIB_RELU_LL.replace("ptr %v", "i8* %v"); + check_implementation_symbols(&legacy, &descriptions).unwrap(); + + // A missing symbol names what the module does define. + let renamed = TILELIB_RELU_LL.replace("@relu_impl(", "@relu("); + let problems = check_implementation_symbols(&renamed, &descriptions).unwrap_err(); + assert_eq!(problems.len(), 1, "{problems:?}"); + assert!( + problems[0].contains("no `define ... @relu_impl(`"), + "{problems:?}" + ); + assert!(problems[0].contains("@scale_in_place_impl"), "{problems:?}"); + + // Every mismatch is reported, symbol by symbol: a wrong scalar type, a + // by-value parameter where a pointer is predicted, a non-void return. + let wrong = TILELIB_RELU_LL + .replace( + "@relu_impl(float %v0, ptr %v1)", + "@relu_impl(double %v0, ptr %v1)", + ) + .replace( + "void @scale_in_place_impl(ptr %v0, float %v1)", + "float @scale_in_place_impl(float %v0, float %v1)", + ); + let problems = check_implementation_symbols(&wrong, &descriptions).unwrap_err(); + assert_eq!(problems.len(), 3, "{problems:?}"); + assert!( + problems[0].contains("`relu_impl`: parameter 0 (`x`) is `double`") + && problems[0].contains("predicts `float`"), + "{problems:?}" + ); + assert!( + problems[1].contains("`scale_in_place_impl`: returns `float`"), + "{problems:?}" + ); + assert!( + problems[2].contains("`scale_in_place_impl`: parameter 0 (`x`) is `float`") + && problems[2].contains("predicts `ptr"), + "{problems:?}" + ); + + // A parameter count mismatch is one problem with both lists. + let short = TILELIB_RELU_LL.replace( + "@warp_sum_impl(float %v0, ptr %v1)", + "@warp_sum_impl(float %v0)", + ); + let problems = check_implementation_symbols(&short, &descriptions).unwrap_err(); + assert_eq!(problems.len(), 1, "{problems:?}"); + assert!( + problems[0].contains("`warp_sum_impl`: 1 parameter(s) in the module (float), 2 in the description (float, ptr"), + "{problems:?}" + ); + + // A definition the module does not export cannot be linked as + // `.extern .func`, however right its signature is. + let hidden = TILELIB_RELU_LL + .replace( + "define void @relu_impl(", + "define internal void @relu_impl(", + ) + .replace( + "define void @warp_sum_impl(", + "define linkonce_odr void @warp_sum_impl(", + ); + let problems = check_implementation_symbols(&hidden, &descriptions).unwrap_err(); + assert_eq!(problems.len(), 2, "{problems:?}"); + assert!( + problems[0].contains("`relu_impl`: defined with `internal` linkage") + && problems[0].contains("not externally visible") + && problems[0].contains(".extern .func"), + "{problems:?}" + ); + assert!( + problems[1].contains("`warp_sum_impl`: defined with `linkonce_odr` linkage"), + "{problems:?}" + ); +} + +/// The capability decision reads `tileiras --list-versions` through +/// the command's own lookup, so the stand-ins here are executables, not +/// listings: that is the seam the command uses. +#[cfg(unix)] +#[test] +fn tileiras_capability_follows_the_listed_bytecode_versions() { + use super::tilelib::list_bytecode_versions; + use oxide_tilebc::BytecodeVersion; + use std::os::unix::fs::PermissionsExt; + + let root = unique_temp_dir("cargo_oxide_tileiras"); + fs::create_dir_all(&root).unwrap(); + let stand_in = |name: &str, body: &str| -> PathBuf { + let path = root.join(name); + fs::write(&path, format!("#!/bin/sh\n{body}\n")).unwrap(); + fs::set_permissions(&path, fs::Permissions::from_mode(0o755)).unwrap(); + path + }; + + // CUDA 13.4's assembler. + let old = stand_in("tileiras-13.4", "printf '13.1\\n13.2\\n13.3\\n13.4\\n'"); + let versions = list_bytecode_versions(&old).unwrap(); + assert_eq!( + versions.iter().map(ToString::to_string).collect::>(), + ["13.1", "13.2", "13.3", "13.4"] + ); + assert!(super::tilelib::description_version(&versions).is_none()); + // The first assembler that reads the preview content; the header tag + // bits are not part of a release. + let new = stand_in( + "tileiras-13.5", + "printf '13.1\\n13.2\\n13.3\\n13.4\\n13.5\\n'", + ); + assert_eq!( + super::tilelib::description_version(&list_bytecode_versions(&new).unwrap()), + Some(BytecodeVersion::PREVIEW_FIRST) + ); + assert_eq!(BytecodeVersion::PREVIEW_FIRST.to_string(), "13.5"); + // The descriptions are written at the assembler's current version, the + // newest numbered one it lists: a newer assembler gets its own version + // (a named build such as `internal` is not a version), a 13.5.1 listing + // names 13.5, and a 13.4 listing names nothing. + let dev = stand_in( + "tileiras-newer", + "printf '13.1\\n13.2\\n13.3\\n13.4\\n13.5\\n13.6\\ninternal\\n'", + ); + assert_eq!( + super::tilelib::description_version(&list_bytecode_versions(&dev).unwrap()) + .map(|version| version.to_string()), + Some("13.6".to_string()) + ); + assert_eq!( + super::tilelib::description_version(&[BytecodeVersion::new(13, 5)]), + Some(BytecodeVersion::PREVIEW_FIRST) + ); + assert_eq!( + super::tilelib::description_version(&list_bytecode_versions(&old).unwrap()), + None + ); + // An assembler without the flag prints a usage message and exits 1: it + // lists nothing, which is not an error. + let legacy = stand_in( + "tileiras-13.1", + "echo 'tileiras: unknown option --list-versions' >&2; exit 1", + ); + assert_eq!(list_bytecode_versions(&legacy).unwrap(), Vec::new()); + assert!(super::tilelib::description_version(&[]).is_none()); + // A configured tileiras that cannot be run is an error naming the path + // and the OS error, not an old assembler. + let missing = root.join("no-such-tileiras"); + let error = list_bytecode_versions(&missing).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::NotFound); + let message = error.to_string(); + assert!( + message.contains(&missing.display().to_string()) + && message.contains("--list-versions") + && message.contains("os error"), + "{message}" + ); + // The exit policy when no library was created. + assert_eq!(no_library_exit_code(true), 0); + assert_eq!(no_library_exit_code(false), 2); + fs::remove_dir_all(root).unwrap(); +} + +#[test] +fn tilelib_paths_and_the_tileiras_invocation_follow_the_crate_stem() { + let dir = Path::new("/tmp/cargo-oxide-tilelib"); + let paths = TileLibraryPaths::new(dir, "my-lib", None); + assert_eq!(paths.ll, dir.join("my_lib.ll")); + assert_eq!(paths.descriptions, dir.join("my_lib.tileexports.json")); + assert_eq!(paths.ltoir, dir.join("my_lib.ltoir")); + assert_eq!(paths.interface, dir.join("my_lib.iface.tilebc")); + assert_eq!(paths.implementation, dir.join("my_lib.impl.tilebc")); + assert_eq!(paths.dump, dir.join("my_lib.tilelib.mlir")); + assert_eq!(paths.library, dir.join("my_lib.tilelib")); + let custom = TileLibraryPaths::new(dir, "my-lib", Some(Path::new("out/lib.tilelib"))); + assert_eq!(custom.library, PathBuf::from("out/lib.tilelib")); + + // The invocation the Tile IR repository documents for creating a + // library: interface, implementation catalog, code payload, output. + assert_eq!( + create_tile_lib_command(Path::new("/opt/cuda/bin/tileiras"), &custom), + [ + "/opt/cuda/bin/tileiras", + "--preview-create-tile-lib", + "--preview-interface=/tmp/cargo-oxide-tilelib/my_lib.iface.tilebc", + "--preview-implementation=/tmp/cargo-oxide-tilelib/my_lib.impl.tilebc", + "--preview-implementation=/tmp/cargo-oxide-tilelib/my_lib.ltoir", + "-o", + "out/lib.tilelib", + ] + ); + // tileiras resolves like every other toolkit tool. + for path in TILEIRAS_FALLBACK_PATHS { + assert_eq!( + Path::new(path).file_name().and_then(|n| n.to_str()), + Some("tileiras") + ); + } +} + +#[test] +fn tilelib_artifacts_are_cleaned_between_builds() { + // A stale library or description must not outlive the build that + // produced it, so every suffix the command writes is a generated + // artifact. + for suffix in [ + TILELIB_INTERFACE_SUFFIX, + TILELIB_IMPLEMENTATION_SUFFIX, + TILELIB_DUMP_SUFFIX, + TILELIB_LIBRARY_SUFFIX, + ] { + assert!(GENERATED_ARTIFACT_SUFFIXES.contains(&suffix), "{suffix}"); + } + let root = unique_temp_dir("cargo_oxide_tilelib_clean"); + std::fs::create_dir_all(&root).unwrap(); + let files: Vec<_> = ["iface.tilebc", "impl.tilebc", "tilelib.mlir", "tilelib"] + .iter() + .map(|suffix| root.join(format!("my_lib.{suffix}"))) + .collect(); + for file in &files { + std::fs::write(file, b"stale").unwrap(); + } + clean_generated_files(&root, "my-lib"); + for file in &files { + assert!(!file.exists(), "{} survived the clean", file.display()); + } + std::fs::remove_dir_all(root).unwrap(); +} + +/// The smoketest's tilelib lane greps the summary lines: the descriptions +/// marker for proof that the descriptions were written before the host +/// program ran, and the two library markers to decide whether a `.tilelib` +/// is required, instead of probing tileiras itself. The script and the +/// command have to name the same text. +#[test] +fn the_smoketest_greps_the_summary_markers_this_command_prints() { + const SMOKETEST: &str = include_str!("../../../../scripts/smoketest.sh"); + for marker in [ + DESCRIPTIONS_WRITTEN_MARKER, + LIBRARY_WRITTEN_MARKER, + LIBRARY_SKIPPED_MARKER, + ] { + assert!( + SMOKETEST.contains(&format!("grep -qF '{marker}'")), + "scripts/smoketest.sh no longer greps {marker:?}" + ); + } + // Each grep must hit only its own line: the library markers must not be + // substrings of the descriptions line, which every run prints. + assert!(!DESCRIPTIONS_WRITTEN_MARKER.contains(LIBRARY_WRITTEN_MARKER)); + assert!(!DESCRIPTIONS_WRITTEN_MARKER.contains(LIBRARY_SKIPPED_MARKER)); + assert!(!LIBRARY_WRITTEN_MARKER.contains(LIBRARY_SKIPPED_MARKER)); +} + /// The project's host CPU flags reach a build only for the triple they name. /// One `cargo oxide` config serves an x86_64 workstation and an aarch64 host /// (a DGX Spark): `target-cpu=znver3` is not an aarch64 processor, and rustc diff --git a/cuda-oxide/crates/cargo-oxide/src/commands/tilelib.rs b/cuda-oxide/crates/cargo-oxide/src/commands/tilelib.rs new file mode 100644 index 0000000000..025d33dbc8 --- /dev/null +++ b/cuda-oxide/crates/cargo-oxide/src/commands/tilelib.rs @@ -0,0 +1,928 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! `cargo oxide tilelib`: package a crate's `#[tile_export]` functions as a +//! cuTile tile library. +//! +//! ```text +//! emit-ltoir ──> .ltoir .ll .tileexports.json +//! │ │ +//! define @(...) checked against abi_params() +//! │ +//! TileExportDesc ──> oxide-tilebc Declaration + Implementation +//! │ +//! .iface.tilebc .impl.tilebc .tilelib.mlir +//! │ +//! tileiras --preview-create-tile-lib (lists 13.5 or newer) ──> .tilelib +//! ``` +//! +//! The descriptions are Tile IR public-preview content, which exists from +//! bytecode 13.5 on and which an assembler reads only at the version it +//! calls current, the newest numbered version its `--list-versions` prints. +//! The command therefore asks the installed tileiras first and writes the +//! descriptions at that version (13.5 for the CUDA 13.5 assembler, and the +//! newer version a newer one names), or at 13.5 when no assembler can read +//! them yet. An older +//! assembler stops the command after the descriptions with an explanation; +//! `--library-optional` makes that outcome a success. Either way the summary +//! ends with one of two lines scripts can read: [`LIBRARY_WRITTEN_MARKER`] +//! or [`LIBRARY_SKIPPED_MARKER`]. + +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use oxide_tile_export::{ + AbiParamKind, DESCRIPTIONS_SUFFIX, TileAccess, TileElem, TileExportDesc, TileSync, +}; +use oxide_tilebc::{ + Access, BytecodeVersion, Declaration, FragmentArg, FragmentType, Implementation, Layout, + Library, MemorySpace, ScalarType, Synchronization, TileArg, TileType, +}; + +use super::*; + +/// cuTile's own override for the assembler, honoured here so one variable +/// selects the same `tileiras` for both tools. +pub(super) const TILEIRAS_PATH_ENV: &str = "CUTILE_TILEIRAS_PATH"; + +/// Suffix of the interface bytecode (`preview$func extern` declarations). +pub(super) const TILELIB_INTERFACE_SUFFIX: &str = "iface.tilebc"; +/// Suffix of the implementation catalog bytecode (`preview$extern_impl`). +pub(super) const TILELIB_IMPLEMENTATION_SUFFIX: &str = "impl.tilebc"; +/// Suffix of the readable dump of both modules. +pub(super) const TILELIB_DUMP_SUFFIX: &str = "tilelib.mlir"; +/// Suffix of the library `tileiras` assembles from the three inputs. +pub(super) const TILELIB_LIBRARY_SUFFIX: &str = "tilelib"; + +/// Fallback locations probed for `tileiras`, in the idiom of the other +/// toolkit tools. It ships in the toolkit's `bin/` beside `nvcc`. +pub(super) const TILEIRAS_FALLBACK_PATHS: &[&str] = &[ + "/usr/local/cuda/bin/tileiras", + "/opt/cuda/bin/tileiras", + "/usr/bin/tileiras", +]; + +/// The line the summary prints once both bytecode halves exist. Scripts +/// grep for it (the smoketest's tilelib lane), so it is a constant here. +pub(super) const DESCRIPTIONS_WRITTEN_MARKER: &str = "Tile library descriptions written"; + +/// The line the summary ends with when tileiras wrote the `.tilelib`; the +/// path and size follow it. The smoketest requires the library when it +/// sees this line, so the script and the command share the text. +pub(super) const LIBRARY_WRITTEN_MARKER: &str = "Tile library written to"; + +/// The line the summary ends with when no library was written, followed by +/// the tileiras that was found and the bytecode versions it lists (or that +/// none was found). The smoketest accepts the lane without a `.tilelib` only +/// when it sees this line. +pub(super) const LIBRARY_SKIPPED_MARKER: &str = "Tile library skipped"; + +// ============================================================================= +// tilelib command +// ============================================================================= + +/// Package a crate's `#[tile_export]` functions as a cuTile tile library. +/// +/// Runs [`emit_ltoir`] for the device half, checks that the emitted NVVM IR +/// defines every implementation symbol with the parameter list the ABI table +/// predicts, maps the descriptions to oxide-tilebc's model, writes the +/// interface and implementation bytecode plus a readable dump, and asks +/// `tileiras --preview-create-tile-lib` for the library when the installed +/// assembler lists bytecode 13.5 or newer, writing the descriptions at its +/// current version. Exits 1 on any error, including a tileiras that cannot +/// be run; exits 2 when no library could be created and `library_optional` +/// was not passed. +#[allow(clippy::too_many_arguments)] +pub fn tilelib( + ctx: &Context, + example: &str, + arch: &str, + features: Option<&str>, + output: Option<&Path>, + library_optional: bool, + print: bool, + verbose: bool, + no_fmad: bool, + unchecked_indexing: bool, + device_debug: DeviceDebug, +) { + let example_dir = if ctx.is_workspace { + resolve_example_dir(ctx, example) + } else { + ctx.workspace_root.clone() + }; + + // Step 1: the device half. Exits on a build or libNVVM failure, and + // clears stale descriptions and libraries before building. + emit_ltoir( + ctx, + example, + arch, + features, + None, + verbose, + no_fmad, + unchecked_indexing, + device_debug, + ); + + let paths = TileLibraryPaths::new(&example_dir, example, output); + let descriptions = read_descriptions(&paths.descriptions, example).unwrap_or_else(|error| { + eprintln!("Error: {error}"); + std::process::exit(1); + }); + + // Step 2: the LTOIR must define what the descriptions promise. The + // producer side of tileiras never looks inside the code payload, so a + // missing or misdeclared symbol would otherwise surface only when a + // cuTile program is compiled against the library. + let ll = std::fs::read_to_string(&paths.ll).unwrap_or_else(|error| { + eprintln!( + "Error: could not read the emitted NVVM IR at {}: {error}", + paths.ll.display() + ); + std::process::exit(1); + }); + if let Err(problems) = check_implementation_symbols(&ll, &descriptions) { + eprintln!( + "Error: {} does not define every implementation the descriptions name:", + paths.ll.display() + ); + for problem in problems { + eprintln!(" {problem}"); + } + std::process::exit(1); + } + + // Step 3: descriptions to Tile IR, verified the way the consumer will. + let stem = artifact_stem(example); + let library = TileLibrary::from_descriptions(&descriptions) + .unwrap_or_else(|error| { + eprintln!("Error: {error}"); + std::process::exit(1); + }) + .library(&stem); + library.verify().unwrap_or_else(|error| { + eprintln!("Error: the tile library does not verify: {error}"); + std::process::exit(1); + }); + + // The assembler decides the bytecode version. Preview content is read + // only at the version an assembler calls current (the newest numbered + // version it lists), so the descriptions are written at that version + // when the installed tileiras can build a library, and at 13.5, the + // first version with preview content, when no assembler can read them + // yet. A tileiras that was named (CUTILE_TILEIRAS_PATH, the toolkit + // root) or found but cannot be started is a configuration error, not an + // old assembler: only a run that lists nothing means "predates + // --list-versions". The error names the path and the OS error. + let tileiras = find_tileiras(ctx); + let versions = match &tileiras { + Some(tileiras) => list_bytecode_versions(tileiras).unwrap_or_else(|error| { + eprintln!("Error: {error}"); + std::process::exit(1); + }), + None => Vec::new(), + }; + let library_version = description_version(&versions); + let version = library_version.unwrap_or(BytecodeVersion::PREVIEW_FIRST); + let interface_bytes = library.interface_bytecode(version).unwrap_or_else(|error| { + eprintln!("Error: could not serialize the interface module: {error}"); + std::process::exit(1); + }); + let implementation_bytes = library + .implementation_bytecode(version) + .unwrap_or_else(|error| { + eprintln!("Error: could not serialize the implementation module: {error}"); + std::process::exit(1); + }); + let dump = library_dump(&library, &stem); + for (path, bytes) in [ + (&paths.interface, interface_bytes.as_slice()), + (&paths.implementation, implementation_bytes.as_slice()), + (&paths.dump, dump.as_bytes()), + ] { + std::fs::write(path, bytes).unwrap_or_else(|error| { + eprintln!("Error: could not write {}: {error}", path.display()); + std::process::exit(1); + }); + } + if print { + println!(); + print!("{dump}"); + } + + println!(); + println!( + "✓ {DESCRIPTIONS_WRITTEN_MARKER}: {} export(s), Tile IR {version} with the preview tag", + descriptions.len(), + ); + println!( + " {} ({} bytes)", + paths.interface.display(), + interface_bytes.len() + ); + println!( + " {} ({} bytes)", + paths.implementation.display(), + implementation_bytes.len() + ); + println!(" {} (readable form of both)", paths.dump.display()); + + // Step 4: the library itself, when the installed assembler can build one. + let (Some(tileiras), Some(version)) = (tileiras.as_deref(), library_version) else { + explain_no_library(&paths, tileiras.as_deref(), &versions); + std::process::exit(no_library_exit_code(library_optional)); + }; + + let command = create_tile_lib_command(tileiras, &paths); + println!(); + println!( + "tileiras {} lists bytecode {version} as its current version; creating the tile library:", + tileiras.display(), + ); + println!(" {}", render_command(&command)); + let output = Command::new(&command[0]) + .args(&command[1..]) + .output() + .unwrap_or_else(|error| { + eprintln!("Error: could not run {}: {error}", tileiras.display()); + std::process::exit(1); + }); + if !output.status.success() { + eprintln!( + "Error: tileiras --preview-create-tile-lib exited with {}", + output.status + ); + eprint!("{}", String::from_utf8_lossy(&output.stderr)); + std::process::exit(1); + } + let size = std::fs::metadata(&paths.library) + .map(|meta| meta.len()) + .unwrap_or(0); + println!(); + println!( + "✓ {LIBRARY_WRITTEN_MARKER} {} ({size} bytes)", + paths.library.display() + ); +} + +/// Every file the command reads or writes for one crate. +pub(super) struct TileLibraryPaths { + /// The NVVM IR the backend emitted, checked for the implementation symbols. + pub(super) ll: PathBuf, + /// The `#[tile_export]` descriptions the backend wrote next to the `.ll`. + pub(super) descriptions: PathBuf, + /// The LTOIR `emit_ltoir` wrote; the code payload of the library. + pub(super) ltoir: PathBuf, + /// The interface bytecode. + pub(super) interface: PathBuf, + /// The implementation catalog bytecode. + pub(super) implementation: PathBuf, + /// The readable dump of both modules. + pub(super) dump: PathBuf, + /// The library `tileiras` writes; `-o` overrides the default. + pub(super) library: PathBuf, +} + +impl TileLibraryPaths { + pub(super) fn new(example_dir: &Path, example: &str, library: Option<&Path>) -> Self { + let stem = artifact_stem(example); + Self { + ll: emitted_ll_path(example_dir, example), + descriptions: example_dir.join(format!("{stem}.{DESCRIPTIONS_SUFFIX}")), + ltoir: default_ltoir_path(example_dir, example), + interface: example_dir.join(format!("{stem}.{TILELIB_INTERFACE_SUFFIX}")), + implementation: example_dir.join(format!("{stem}.{TILELIB_IMPLEMENTATION_SUFFIX}")), + dump: example_dir.join(format!("{stem}.{TILELIB_DUMP_SUFFIX}")), + library: library + .map(Path::to_path_buf) + .unwrap_or_else(|| example_dir.join(format!("{stem}.{TILELIB_LIBRARY_SUFFIX}"))), + } + } +} + +/// The descriptions the backend wrote, validated. A crate without +/// `#[tile_export]` functions has no descriptions file at all, which is the +/// common mistake and gets the direct message. +fn read_descriptions(path: &Path, example: &str) -> Result, String> { + let text = match std::fs::read_to_string(path) { + Ok(text) => text, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + return Err(format!( + "no #[tile_export] functions in {example}: the build wrote no {} (a tile \ + library needs at least one exported function)", + path.display() + )); + } + Err(error) => return Err(format!("could not read {}: {error}", path.display())), + }; + let descriptions = oxide_tile_export::descriptions_from_json(&text) + .map_err(|error| format!("could not parse {}: {error}", path.display()))?; + if descriptions.is_empty() { + return Err(format!( + "no #[tile_export] functions in {example}: {} describes no exports", + path.display() + )); + } + for description in &descriptions { + description.validate().map_err(|error| { + format!( + "description of `{}` in {}: {error}", + description.decl, + path.display() + ) + })?; + } + Ok(descriptions) +} + +// ============================================================================= +// Descriptions to Tile IR +// ============================================================================= + +/// A crate's exports as the two halves of a tile library. +#[derive(Debug)] +pub(super) struct TileLibrary { + /// One `preview$func extern` declaration per export. + pub(super) funcs: Vec, + /// One `preview$extern_impl` per export. + pub(super) impls: Vec, +} + +impl TileLibrary { + /// Map every description; the first one that cannot be described stops + /// the command with its reason. + pub(super) fn from_descriptions(descriptions: &[TileExportDesc]) -> Result { + let mut funcs = Vec::with_capacity(descriptions.len()); + let mut impls = Vec::with_capacity(descriptions.len()); + for description in descriptions { + let (func, implementation) = describe_export(description) + .map_err(|error| format!("export `{}`: {error}", description.decl))?; + funcs.push(func); + impls.push(implementation); + } + Ok(Self { funcs, impls }) + } + + /// Both halves under one module name, ready to verify and write. + pub(super) fn library(&self, name: &str) -> Library { + Library::new(name, self.funcs.clone(), self.impls.clone()) + } +} + +/// One export as its declaration and its implementation. +/// +/// Every parameter becomes a tile of the export's shape on the declaration +/// and a fragment on the implementation: a by-value scalar is the rank-0 +/// `():()` buffer, everything else a contiguous footprint `(F):(1)`. The +/// fragmap is shared by all parameters, see [`fragmap`]. A `bf16` scalar +/// passed by value needs the `bf16 = i16` frontend ABI entry because +/// cuda-device's bf16 is a `u16` wrapper, so the LLVM parameter is `i16` +/// while the tile element is `bf16`. +pub(super) fn describe_export( + desc: &TileExportDesc, +) -> Result<(Declaration, Implementation), String> { + desc.validate().map_err(|error| error.to_string())?; + let tile_size = desc.shape[0]; + // Every parameter is a fragment of a tile of the same shape and the + // fragmaps of one implementation share their invocation count, so the + // footprints have to agree: a 4-element window beside a scalar would + // need 32 invocations for one argument and 128 for the other. + let footprint = desc.params[0].footprint; + if let Some(other) = desc + .params + .iter() + .find(|param| param.footprint != footprint) + { + return Err(format!( + "parameters `{}` (footprint {}) and `{}` (footprint {}) differ; every parameter \ + of one export shares the invocation count, so all footprints must be equal", + desc.params[0].name, footprint, other.name, other.footprint + )); + } + let mut args = Vec::with_capacity(desc.params.len()); + let mut fragments = Vec::with_capacity(desc.params.len()); + for param in &desc.params { + let element = scalar_type(param.elem); + let access = access(param.access); + args.push(TileArg::new( + param.name.clone(), + TileType::new(desc.shape.clone(), element), + access, + )); + let footprint = i64::from(param.footprint); + let fragmap = fragmap(desc.sync, tile_size, footprint) + .map_err(|error| format!("parameter `{}`: {error}", param.name))?; + let (shape, mem_layout) = if param.by_value { + (vec![1], Layout::scalar_mem()) + } else { + (vec![footprint], Layout::contiguous_mem(footprint)) + }; + fragments.push(FragmentArg::new( + param.name.clone(), + FragmentType::new(shape, element, MemorySpace::Local, mem_layout, fragmap), + access, + )); + } + let mut implementation = Implementation::new(desc.symbol.clone(), desc.decl.clone(), fragments); + implementation.synchronization = synchronization(desc.sync); + implementation.arch = desc.arch.clone(); + if desc + .params + .iter() + .any(|param| param.by_value && param.elem == TileElem::BF16) + { + implementation.frontend_abi = vec![("bf16".to_string(), "i16".to_string())]; + } + Ok((Declaration::new(desc.decl.clone(), args), implementation)) +} + +/// How one invocation's footprint maps onto the tile. +/// +/// Without synchronization every invocation owns `footprint` contiguous +/// elements: `(N, 1, 1):(1, 0, 0)` for a footprint of one, `(N/F, 1, F):(F, +/// 0, 1)` otherwise. With `sync = warp` the 32 lanes of a warp are the +/// participants of one invocation: `(N/32, 32, 1):(32, 1, 0)` for a footprint +/// of one, and `(N/(32F), 32, F):(32F, F, 1)` for a footprint of `F`, so each +/// lane keeps a contiguous footprint and adjacent lanes are adjacent in the +/// tile. The footprint must divide the tile size in either case. +pub(super) fn fragmap(sync: TileSync, tile_size: i64, footprint: i64) -> Result { + if footprint <= 0 { + return Err(format!("footprint {footprint} must be positive")); + } + if tile_size % footprint != 0 { + return Err(format!( + "footprint {footprint} does not divide the tile size {tile_size}" + )); + } + match sync { + TileSync::None if footprint == 1 => Ok(Layout::elementwise(tile_size)), + TileSync::None => Ok(Layout::footprint_blocks(tile_size, footprint)), + TileSync::Warp => { + let per_warp = 32 * footprint; + if tile_size % per_warp != 0 { + return Err(format!( + "sync = warp with footprint {footprint} needs a tile size that is a \ + multiple of {per_warp}, found {tile_size}" + )); + } + if footprint == 1 { + Ok(Layout::warp_lanes(tile_size)) + } else { + Ok(Layout::flat( + &[tile_size / per_warp, 32, footprint], + &[per_warp, footprint, 1], + )) + } + } + } +} + +/// Tile IR has no unsigned integers; the unsigned Rust types are the signless +/// `iN` of the same width. +fn scalar_type(elem: TileElem) -> ScalarType { + match elem { + TileElem::F16 => ScalarType::F16, + TileElem::BF16 => ScalarType::BF16, + TileElem::F32 => ScalarType::F32, + TileElem::F64 => ScalarType::F64, + TileElem::I8 | TileElem::U8 => ScalarType::I8, + TileElem::I16 | TileElem::U16 => ScalarType::I16, + TileElem::I32 | TileElem::U32 => ScalarType::I32, + TileElem::I64 | TileElem::U64 => ScalarType::I64, + } +} + +fn access(access: TileAccess) -> Access { + match access { + TileAccess::In => Access::In, + TileAccess::Out => Access::Out, + TileAccess::InOut => Access::InOut, + } +} + +/// The implementation's `synchronization` clause. An export without `sync` +/// gets no clause at all rather than an explicit `synchronization`: +/// the two mean the same thing, and leaving it out keeps the printed and +/// written form identical to an upstream module parsed without the clause. +fn synchronization(sync: TileSync) -> Option { + match sync { + TileSync::None => None, + TileSync::Warp => Some(Synchronization::Warp), + } +} + +/// The readable form of both modules, one dump so a reader sees the +/// declaration and its implementation together. Lines starting with `//` +/// are comments in the Tile IR textual syntax. +pub(super) fn library_dump(library: &Library, stem: &str) -> String { + format!( + "// {stem}.{TILELIB_INTERFACE_SUFFIX}: the interface a cuTile program calls\n{}\ + // {stem}.{TILELIB_IMPLEMENTATION_SUFFIX}: the implementations the LTOIR defines\n{}", + library.interface_text(), + library.implementation_text() + ) +} + +/// The bytecode versions `tileiras --list-versions` names, one per line. +/// +/// A launch failure is returned as the error it is, with the path in the +/// message. An assembler that ran but predates the flag exits non-zero or +/// prints a usage message instead of versions; both give an empty list. +/// Lines that are not versions (an assembler may list a named build next +/// to the numbers) are skipped. +pub(super) fn list_bytecode_versions(tileiras: &Path) -> std::io::Result> { + let output = Command::new(tileiras) + .arg("--list-versions") + .output() + .map_err(|error| { + std::io::Error::new( + error.kind(), + format!( + "could not run {} --list-versions: {error}", + tileiras.display() + ), + ) + })?; + if !output.status.success() { + return Ok(Vec::new()); + } + Ok(String::from_utf8_lossy(&output.stdout) + .lines() + .filter_map(BytecodeVersion::parse) + .collect()) +} + +// ============================================================================= +// Implementation symbols in the emitted NVVM IR +// ============================================================================= + +/// One `define` of an LLVM IR module, reduced to what the ABI check reads. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct LlDefine { + /// The return type, with linkage and calling-convention keywords removed. + pub(super) ret: String, + /// One type per parameter, names and parameter attributes removed. + pub(super) params: Vec, + /// A linkage keyword from [`LL_LOCAL_LINKAGE`] on the define, if any. + pub(super) local_linkage: Option, +} + +/// Linkages that keep a definition out of what the module exports. The Tile +/// compiler reaches an implementation as `.extern .func`, so a definition +/// with one of these is a reported problem, not a stripped keyword. +const LL_LOCAL_LINKAGE: &[&str] = &["private", "internal", "linkonce", "linkonce_odr"]; + +/// Keywords that can sit between `define` and the return type without +/// changing whether the symbol can be linked. +const LL_DEFINE_KEYWORDS: &[&str] = &[ + "available_externally", + "weak", + "common", + "appending", + "extern_weak", + "weak_odr", + "external", + "dso_local", + "dso_preemptable", + "default", + "hidden", + "protected", + "ptx_kernel", + "ptx_device", + "fastcc", + "coldcc", + "ccc", +]; + +/// Every `define ... @name(params)` of `ll`, keyed by symbol. Textual IR is +/// the robust place to check symbols: the LTOIR is an NVVM bitcode container +/// with no stable public reader, and the `.ll` beside it is what libNVVM +/// compiled. +pub(super) fn ll_defines(ll: &str) -> BTreeMap { + let mut defines = BTreeMap::new(); + for line in ll.lines() { + let Some(rest) = line.trim_start().strip_prefix("define ") else { + continue; + }; + let Some(at) = rest.find('@') else { + continue; + }; + let (prefix, after_at) = rest.split_at(at); + let after_at = &after_at[1..]; + let Some(open) = after_at.find('(') else { + continue; + }; + let name = after_at[..open].trim().trim_matches('"').to_string(); + let Some(params) = balanced_parameter_list(&after_at[open..]) else { + continue; + }; + let local_linkage = prefix + .split_whitespace() + .find(|token| LL_LOCAL_LINKAGE.contains(token)) + .map(str::to_string); + let ret = prefix + .split_whitespace() + .filter(|token| { + !LL_DEFINE_KEYWORDS.contains(token) && !LL_LOCAL_LINKAGE.contains(token) + }) + .collect::>() + .join(" "); + defines.insert( + name, + LlDefine { + ret, + params: split_top_level(params) + .into_iter() + .map(parameter_type) + .collect(), + local_linkage, + }, + ); + } + defines +} + +/// The text inside the parentheses that open `text`, honoring nesting. +fn balanced_parameter_list(text: &str) -> Option<&str> { + let mut depth = 0i32; + for (index, c) in text.char_indices() { + match c { + '(' | '[' | '{' | '<' => depth += 1, + ')' | ']' | '}' | '>' => { + depth -= 1; + if depth == 0 { + return Some(&text[1..index]); + } + } + _ => {} + } + } + None +} + +/// Split on commas outside any bracket pair. +fn split_top_level(text: &str) -> Vec<&str> { + let mut parts = Vec::new(); + let mut depth = 0i32; + let mut start = 0; + for (index, c) in text.char_indices() { + match c { + '(' | '[' | '{' | '<' => depth += 1, + ')' | ']' | '}' | '>' => depth -= 1, + ',' if depth == 0 => { + parts.push(&text[start..index]); + start = index + 1; + } + _ => {} + } + } + let last = &text[start..]; + if !last.trim().is_empty() { + parts.push(last); + } + parts +} + +/// The type of one parameter: its tokens without the `%name` and without +/// parameter attributes. +fn parameter_type(param: &str) -> String { + const ATTRIBUTES: &[&str] = &[ + "noundef", + "nonnull", + "noalias", + "nocapture", + "readonly", + "writeonly", + "readnone", + "signext", + "zeroext", + "inreg", + "returned", + "nofree", + "nest", + "immarg", + ]; + param + .split_whitespace() + .filter(|token| !token.starts_with('%')) + .filter(|token| !ATTRIBUTES.contains(token)) + .filter(|token| !token.starts_with("align") && !token.starts_with("dereferenceable")) + .collect::>() + .join(" ") +} + +/// Whether an LLVM parameter type is the one the ABI predicts. A pointer is +/// the opaque `ptr` on the modern NVVM path and a typed pointer (`i8*`) in +/// the legacy LLVM 7 dialect of pre-Blackwell targets; the pointee type +/// carries no ABI information, so any typed pointer passes. +fn llvm_type_matches(actual: &str, kind: AbiParamKind) -> bool { + match kind { + AbiParamKind::Scalar(elem) => actual == elem.llvm_name(), + AbiParamKind::Pointer(_) => actual == "ptr" || actual.ends_with('*'), + } +} + +/// The type text the message quotes for an expected parameter. +fn expected_type_text(kind: AbiParamKind) -> String { + match kind { + AbiParamKind::Scalar(elem) => elem.llvm_name().to_string(), + AbiParamKind::Pointer(_) => "ptr (i8* on the legacy path)".to_string(), + } +} + +/// Check that `ll` defines every description's implementation symbol as a +/// `void` function whose parameters are the ones [`TileExportDesc::abi_params`] +/// predicts. Every mismatch is reported, symbol by symbol. +pub(super) fn check_implementation_symbols( + ll: &str, + descriptions: &[TileExportDesc], +) -> Result<(), Vec> { + let defines = ll_defines(ll); + let mut problems = Vec::new(); + for desc in descriptions { + let Some(define) = defines.get(&desc.symbol) else { + let defined: Vec<_> = defines.keys().map(|name| format!("@{name}")).collect(); + problems.push(format!( + "`{}` (implementation of `{}`): no `define ... @{}(` in the module; it defines {}", + desc.symbol, + desc.decl, + desc.symbol, + if defined.is_empty() { + "no functions".to_string() + } else { + defined.join(", ") + } + )); + continue; + }; + if let Some(linkage) = &define.local_linkage { + problems.push(format!( + "`{}`: defined with `{linkage}` linkage, so it is not externally visible; \ + the Tile compiler links it as .extern .func", + desc.symbol + )); + } + if define.ret != "void" { + problems.push(format!( + "`{}`: returns `{}`, but a tile library implementation returns void and \ + writes its result through the trailing out pointer", + desc.symbol, define.ret + )); + } + let expected = desc.abi_params(); + if define.params.len() != expected.len() { + problems.push(format!( + "`{}`: {} parameter(s) in the module ({}), {} in the description ({})", + desc.symbol, + define.params.len(), + define.params.join(", "), + expected.len(), + expected + .iter() + .map(|param| expected_type_text(param.kind)) + .collect::>() + .join(", ") + )); + continue; + } + for (index, (actual, param)) in define.params.iter().zip(&expected).enumerate() { + if !llvm_type_matches(actual, param.kind) { + problems.push(format!( + "`{}`: parameter {index} (`{}`) is `{actual}` in the module, but the \ + description predicts `{}`", + desc.symbol, + param.name, + expected_type_text(param.kind) + )); + } + } + } + if problems.is_empty() { + Ok(()) + } else { + Err(problems) + } +} + +// ============================================================================= +// tileiras +// ============================================================================= + +/// `tileiras` from cuTile's override variable (also honoured as +/// [`TILEIRAS_PATH_ENV`], so one variable selects the same assembler for +/// both products), then the toolkit lookup every other CUDA tool uses +/// (PATH, the configured toolkit root, the standard install roots). An +/// override is returned as given: whether it runs is for +/// [`list_bytecode_versions`] to report. +pub(super) fn find_tileiras(ctx: &Context) -> Option { + if let Some(path) = std::env::var_os(TILEIRAS_PATH_ENV).filter(|value| !value.is_empty()) { + return Some(PathBuf::from(path)); + } + find_cuda_toolkit_executable(ctx, "tileiras", TILEIRAS_FALLBACK_PATHS) +} + +/// The bytecode version to write the descriptions at for an assembler that +/// lists `versions`: its current version, the newest numbered one it lists, +/// when that is 13.5 or newer. An assembler reads preview content only at +/// its current version, so an older listing (CUDA 13.4 and before) means no +/// library can be built. The header tag bits are not part of a release, so +/// `13.5` and `13.5.1` both name 13.5. +pub(super) fn description_version(versions: &[BytecodeVersion]) -> Option { + versions + .iter() + .copied() + .filter(|version| *version >= BytecodeVersion::PREVIEW_FIRST) + .max() +} + +/// Exit status when no library was created because the installed assembler +/// cannot build one: a success when the caller said the library is optional, +/// and a shortfall otherwise. +pub(super) fn no_library_exit_code(library_optional: bool) -> i32 { + if library_optional { 0 } else { 2 } +} + +/// The `tileiras --preview-create-tile-lib` invocation, argv form. +pub(super) fn create_tile_lib_command(tileiras: &Path, paths: &TileLibraryPaths) -> Vec { + vec![ + tileiras.display().to_string(), + "--preview-create-tile-lib".to_string(), + format!("--preview-interface={}", paths.interface.display()), + format!( + "--preview-implementation={}", + paths.implementation.display() + ), + format!("--preview-implementation={}", paths.ltoir.display()), + "-o".to_string(), + paths.library.display().to_string(), + ] +} + +fn render_command(command: &[String]) -> String { + command.join(" ") +} + +/// The [`LIBRARY_SKIPPED_MARKER`] line naming the assembler that was found +/// and what it lists, then why that is not enough and everything a reader +/// needs to finish the job on a newer toolkit. +fn explain_no_library( + paths: &TileLibraryPaths, + tileiras: Option<&Path>, + versions: &[BytecodeVersion], +) { + let needed = BytecodeVersion::PREVIEW_FIRST; + println!(); + match tileiras { + Some(tileiras) => { + let lists = if versions.is_empty() { + "no bytecode versions (it predates --list-versions)".to_string() + } else { + format!( + "bytecode versions {}", + versions + .iter() + .map(ToString::to_string) + .collect::>() + .join(", ") + ) + }; + println!( + "{LIBRARY_SKIPPED_MARKER}: {} lists {lists}; the descriptions need {needed} or newer", + tileiras.display() + ); + println!( + "The installed tileiras cannot create tile libraries: the descriptions are \ + Tile IR public-preview content, which exists from bytecode {needed} on and \ + needs the CUDA 13.5 or newer assembler." + ); + } + None => { + println!("{LIBRARY_SKIPPED_MARKER}: no tileiras found"); + println!( + "No tileiras was found on PATH, under {TILEIRAS_PATH_ENV}, the configured CUDA \ + toolkit root or the standard install roots; creating a tile library needs the \ + CUDA 13.5 or newer assembler, the first one that lists bytecode {needed}." + ); + } + } + println!( + "The LTOIR and the descriptions were written to {}; a capable tileiras turns them \ + into {} with:", + paths + .interface + .parent() + .unwrap_or_else(|| Path::new(".")) + .display(), + paths.library.display() + ); + let placeholder = Path::new("tileiras"); + println!( + " {}", + render_command(&create_tile_lib_command(placeholder, paths)) + ); +} 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 292bed15cd..bfd4524b61 100644 --- a/cuda-oxide/crates/cargo-oxide/src/main.rs +++ b/cuda-oxide/crates/cargo-oxide/src/main.rs @@ -19,6 +19,7 @@ //! cargo oxide debug vecadd --tui # build + cuda-gdb //! cargo oxide inspect vecadd # build + print generated PTX //! cargo oxide fuzz-schedule barrier # perturb generated PTX and watchdog it +//! cargo oxide tilelib relu --arch sm_120 # #[tile_export] crate -> cuTile tile library //! cargo oxide new my_kernel # scaffold a standalone project //! cargo oxide new my_kernel --async # scaffold with async template //! cargo oxide list # list bundled examples @@ -28,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}; @@ -37,6 +39,7 @@ mod artifact_identity; mod backend; mod backend_source; mod commands; +mod cutlass_toolchain; /// Top-level CLI structure parsed by clap. /// @@ -72,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)] @@ -344,6 +347,63 @@ enum Commands { #[arg(long)] unchecked_indexing: bool, }, + // Hidden until the CUDA 13.5 toolkit is announced: the feature needs that + // assembler and gets documented when it ships. + #[command(hide = true)] + /// Package a crate's `#[tile_export]` functions as a cuTile tile library. + /// + /// Runs `emit-ltoir`, checks that the LTOIR defines every implementation + /// symbol with the parameters the native SIMT ABI predicts, writes the + /// Tile IR 13.5 interface and implementation bytecode + /// (`.iface.tilebc`, `.impl.tilebc`) plus a readable + /// `.tilelib.mlir`, and asks `tileiras --preview-create-tile-lib` + /// for `.tilelib` when the installed assembler lists bytecode 13.5 + /// (CUDA 13.5 or newer). An older assembler stops the command after the + /// descriptions with exit code 2 unless `--library-optional` is passed. + Tilelib { + /// Crate name (required in workspace, optional for standalone projects) + example: Option, + /// Target architecture (e.g. sm_90, sm_100, sm_120). Required: the + /// LTOIR inside the library is architecture-specific. + #[arg(long)] + arch: String, + /// Comma-separated list of features to enable + #[arg(long)] + features: Option, + /// Output path for the `.tilelib` file (default: `/.tilelib`) + #[arg(long, short)] + output: Option, + /// Succeed when the installed tileiras cannot build the .tilelib; + /// the descriptions and LTOIR are still written (without this flag + /// that case exits 2) + #[arg(long)] + library_optional: bool, + /// Print the readable form of the interface and implementation modules + #[arg(long)] + print: bool, + /// Show verbose compilation output + #[arg(short, long)] + verbose: bool, + /// Disable implicit FMA contraction in both libNVVM and nvJitLink. + /// Also settable via CUDA_OXIDE_NO_FMA=1. + #[arg(long)] + no_fmad: bool, + /// Emit device line-number information for profilers and Compute + /// Sanitizer, leaving optimization intact (nvcc `-lineinfo`). + /// Also settable via CUDA_OXIDE_DEBUG=line. + #[arg(long)] + lineinfo: bool, + /// Emit full device debug information; libNVVM finalization runs + /// unoptimized (nvcc `-G`). Supersedes --lineinfo. + /// Also settable via CUDA_OXIDE_DEBUG=full. + #[arg(long)] + device_debug: bool, + /// Elide slice/array bounds checks in every device kernel + /// (out-of-bounds indexing becomes UB, like get_unchecked). + /// Also settable via CUDA_OXIDE_UNCHECKED_INDEXING=1. + #[arg(long)] + unchecked_indexing: bool, + }, /// Show the full compilation pipeline (MIR -> PTX/NVVM IR) with verbose output Pipeline { /// Example name (required in workspace, optional for standalone projects) @@ -397,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)] @@ -462,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)] @@ -585,6 +665,10 @@ fn validate_materialization_cli(cli: &Cli) -> Result<(), String> { "--materialize-cubin cannot be combined with emit-ltoir; one emits a final cubin and the other emits linkable LTOIR" .to_string(), ), + Commands::Tilelib { .. } => Err( + "--materialize-cubin cannot be combined with tilelib; a tile library carries linkable LTOIR, not a final cubin" + .to_string(), + ), Commands::List { .. } => Err( "--materialize-cubin cannot be used with list because list does not compile device code" .to_string(), @@ -613,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(), @@ -985,6 +1073,35 @@ fn main() { commands::DeviceDebug::from_flags(lineinfo, device_debug), ); } + Commands::Tilelib { + example, + arch, + features, + output, + library_optional, + print, + verbose, + no_fmad, + unchecked_indexing, + lineinfo, + device_debug, + } => { + let ctx = commands::resolve_context(); + let example = resolve_example_name(example, &ctx, "tilelib"); + commands::tilelib( + &ctx, + &example, + &arch, + features.as_deref(), + output.as_deref(), + library_optional, + print, + verbose, + no_fmad, + unchecked_indexing, + commands::DeviceDebug::from_flags(lineinfo, device_debug), + ); + } Commands::Pipeline { example, emit_nvvm_ir, @@ -1099,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); + } + }, + }, } } @@ -1124,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); } @@ -1215,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(&[ @@ -1326,6 +1470,14 @@ mod tests { "--materialize-cubin", ], &["cargo-oxide", "inspect", "demo", "--materialize-cubin"], + &[ + "cargo-oxide", + "tilelib", + "demo", + "--arch", + "sm_120", + "--materialize-cubin", + ], ] { let cli = Cli::try_parse_from(args).expect("global flag should parse first"); let error = validate_materialization_cli(&cli) @@ -1679,6 +1831,65 @@ mod tests { assert!(!device_debug, "--device-debug must default off"); } + #[test] + fn tilelib_parser_requires_arch_and_accepts_its_flags() { + assert!( + Cli::try_parse_from(["cargo-oxide", "tilelib", "relu"]).is_err(), + "the LTOIR inside a tile library is architecture-specific, so --arch is required" + ); + + let cli = Cli::try_parse_from([ + "cargo-oxide", + "tilelib", + "relu", + "--arch", + "sm_120", + "--library-optional", + "--print", + "-o", + "out/relu.tilelib", + ]) + .expect("tilelib command should parse"); + let Commands::Tilelib { + example, + arch, + output, + library_optional, + print, + verbose, + no_fmad, + lineinfo, + device_debug, + .. + } = cli.command + else { + panic!("expected tilelib command"); + }; + assert_eq!(example.as_deref(), Some("relu")); + assert_eq!(arch, "sm_120"); + assert_eq!(output, Some(PathBuf::from("out/relu.tilelib"))); + assert!(library_optional); + assert!(print); + assert!(!verbose); + assert!(!no_fmad); + assert!(!lineinfo, "--lineinfo must default off"); + assert!(!device_debug, "--device-debug must default off"); + + // Without the two flags the command insists on a library. + let cli = Cli::try_parse_from(["cargo-oxide", "tilelib", "relu", "--arch", "sm_120"]) + .expect("tilelib command should parse"); + let Commands::Tilelib { + library_optional, + print, + .. + } = cli.command + else { + panic!("expected tilelib command"); + }; + assert!(!library_optional); + assert!(!print); + } + #[test] fn parser_accepts_device_debug_flags() { let cli = Cli::try_parse_from(["cargo-oxide", "build", "vecadd", "--lineinfo"]) diff --git a/cuda-oxide/crates/cuda-device/README.md b/cuda-oxide/crates/cuda-device/README.md index 9b43073eb6..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 │ @@ -184,16 +188,15 @@ pub fn debug_kernel(data: &[f32]) { let idx = thread::index_1d(); let value = data[idx.get()]; - gpu_printf!( - "thread %d: val = %f\n", - idx.get() as i32, - value as f64 - ); + gpu_printf!("thread {}: val = {}\n", idx.get(), value); gpu_assert!(value >= 0.0, "expected non-negative value"); } ``` -`gpu_printf!` compiles to device-side `vprintf` with C vararg promotion. +`gpu_printf!` compiles to device-side `vprintf` with C vararg promotion. A +placeholder without a type character prints its argument by type (`%llu` for +the `usize` index, `%f` for the `f32` value); `{:x}`, `{:.2f}` and the other C +type characters cast the argument instead. `gpu_assert!(condition)` uses `trap()` when the condition is false. The string-literal message form, `gpu_assert!(condition, "message")`, calls CUDA's @@ -223,6 +226,7 @@ convenience. The table is the whole `pub use cuda_macros::{...}` list in | `#[pure]` | Mark as pure (no side effects) | | `#[readonly]` | Mark as read-only | | `#[constant]` | Place a `ConstantMemory` static in constant memory, plus a host `set_` setter | +| `#[tile_export]` | Export a device function to cuTile through a tile library | | `gpu_printf!` | Device-side printf | | `ptx_asm!` | Unsafe CUDA inline PTX | diff --git a/cuda-oxide/crates/cuda-device/src/atomic.rs b/cuda-oxide/crates/cuda-device/src/atomic.rs index c39c61cd25..fa616afd12 100644 --- a/cuda-oxide/crates/cuda-device/src/atomic.rs +++ b/cuda-oxide/crates/cuda-device/src/atomic.rs @@ -201,8 +201,10 @@ macro_rules! define_integer_atomic { /// atomic and non-atomic access to the same location is UB. /// * For block-scoped types, only threads in the same thread block /// may access this memory. - /// * For system-scoped types, the memory must be reachable by both - /// the GPU and the CPU (unified memory). + /// * The memory and concurrently accessing threads must support + /// the chosen scope. System scope permits device memory accessed + /// only by threads on one GPU; CPU or peer-GPU access must satisfy + /// the [CUDA atomicity requirements](https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cuda-cpp-memory-model.html#atomicity). #[inline(always)] pub const unsafe fn from_ptr<'a>(ptr: *mut $ty) -> &'a Self { // SAFETY: `Self` is `#[repr(transparent)]` over `UnsafeCell<$ty>`, @@ -388,7 +390,9 @@ macro_rules! define_float_atomic { /// the entire lifetime `'a`. /// * No non-atomic accesses to `*ptr` may occur during `'a`. /// * Block-scoped: only same-block threads may access. - /// System-scoped: memory must be reachable by GPU and CPU. + /// * The memory and concurrently accessing threads must support + /// the chosen scope, as described on the integer atomic types. + /// System scope permits device memory accessed only by one GPU. #[inline(always)] pub const unsafe fn from_ptr<'a>(ptr: *mut $ty) -> &'a Self { // SAFETY: `Self` is `#[repr(transparent)]` over `UnsafeCell<$ty>`, @@ -490,7 +494,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-08-28` 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 703213caf3..82a8934c0c 100644 --- a/cuda-oxide/crates/cuda-device/src/barrier.rs +++ b/cuda-oxide/crates/cuda-device/src/barrier.rs @@ -254,29 +254,29 @@ impl BarrierToken { /// /// # Thread Requirements /// -/// | Operation | Thread Requirement | -/// |-----------------|---------------------------| -/// | `from_static()` | Single thread (thread 0) | -/// | `init()` | Single thread (thread 0) | -/// | `arrive()` | All participating threads | -/// | `wait()` | All participating threads | -/// | `inval()` | Single thread (thread 0) | +/// | Operation | Thread Requirement | +/// |-----------------|----------------------------------------| +/// | `from_static()` | All threads call | +/// | `init()` | All threads call; linear thread 0 acts | +/// | `arrive()` | All participating threads | +/// | `wait()` | All participating threads | +/// | `inval()` | All threads call; linear thread 0 acts | /// /// # Example /// /// ```rust,ignore -/// // Thread 0: Create and initialize +/// // All threads: create the handle and initialize. Linear thread 0, that is +/// // thread (0, 0, 0), runs mbarrier.init; init() then fences and +/// // synchronizes the block, so every thread holds a Ready handle. /// let bar = ManagedBarrier::::from_static(&raw mut BAR); /// let bar = unsafe { bar.init(128) }; // Now Ready -/// fence_proxy_async_shared_cta(); -/// sync_threads(); /// /// // All threads: arrive and wait /// let token = bar.arrive(); /// bar.wait(token); /// -/// // Thread 0: Cleanup -/// sync_threads(); +/// // All threads: clean up. inval() synchronizes the block, then linear +/// // thread 0 runs mbarrier.inval. /// unsafe { bar.inval(); } /// ``` pub struct ManagedBarrier { @@ -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 @@ -301,7 +316,11 @@ impl ManagedBarrier { /// /// # Thread Requirements /// - /// **Single-thread**: Only call from the initialization thread (thread 0). + /// **All threads**: this is a plain constructor with no device effect. Every + /// thread wraps the same static, and `init()` then elects linear thread 0 to + /// run `mbarrier.init`. Do not wrap `init()` in a thread-0 branch: it + /// synchronizes the whole block, and a block-wide sync inside a branch that + /// only some threads take is undefined, in practice a hang. /// /// # Example /// @@ -309,13 +328,12 @@ impl ManagedBarrier { /// static mut BAR0: Barrier = Barrier::UNINIT; /// static mut BAR1: Barrier = Barrier::UNINIT; /// - /// if tid == 0 { - /// let bar0 = ManagedBarrier::::from_static(&raw mut BAR0); - /// let bar1 = ManagedBarrier::::from_static(&raw mut BAR1); + /// // All threads: wrap and initialize; init() elects linear thread 0. + /// let bar0 = ManagedBarrier::::from_static(&raw mut BAR0); + /// let bar1 = ManagedBarrier::::from_static(&raw mut BAR1); /// - /// let bar0 = unsafe { bar0.init(32) }; - /// let bar1 = unsafe { bar1.init(32) }; - /// } + /// let bar0 = unsafe { bar0.init(32) }; + /// let bar1 = unsafe { bar1.init(32) }; /// ``` pub fn from_static(ptr: *mut Barrier) -> Self { ManagedBarrier { @@ -327,8 +345,9 @@ impl ManagedBarrier { /// Initialize the barrier with an expected arrival count. /// - /// **All threads in the block should call this.** Only thread 0 performs - /// the actual initialization; all threads synchronize and receive a `Ready` handle. + /// **All threads in the block should call this.** Only linear thread 0, + /// that is thread (0, 0, 0), performs the actual initialization; all threads + /// synchronize and receive a `Ready` handle. /// /// This is a convenience wrapper for `init_by(count, 0)`. /// @@ -347,7 +366,7 @@ impl ManagedBarrier { /// ```rust,ignore /// static mut BAR: Barrier = Barrier::UNINIT; /// - /// // ALL threads call init - only thread 0 actually initializes + /// // ALL threads call init - only linear thread 0 actually initializes /// let bar = ManagedBarrier::::from_static(&raw mut BAR); /// let bar = unsafe { bar.init(128) }; // All threads get Ready handle /// ``` @@ -358,8 +377,8 @@ impl ManagedBarrier { /// Initialize the barrier with a specific thread performing initialization. /// - /// **All threads in the block should call this.** Only the thread with - /// `threadIdx.x == init_thread` performs the actual initialization; + /// **All threads in the block should call this.** Only the thread whose + /// linear id equals `init_thread` performs the actual initialization; /// all threads synchronize and receive a `Ready` handle. /// /// # Block-Scoped Barriers @@ -371,13 +390,18 @@ impl ManagedBarrier { /// # Parameters /// /// - `count`: Expected number of arrivals before barrier completes - /// - `init_thread`: Thread ID (threadIdx.x) that performs initialization + /// - `init_thread`: Linear thread id within the block that performs + /// initialization. `0` is thread `(0, 0, 0)`. In a 1D block this is + /// `threadIdx.x`. /// /// # Safety /// /// - Must be called before any arrive/wait operations /// - All participating threads in the block must call this together - /// - `init_thread` must be a valid thread ID within the block + /// - `init_thread` must be the same value in every thread and below + /// `blockDim.x * blockDim.y * blockDim.z`. A selector that differs between + /// threads elects zero or several threads; an out-of-range id selects no + /// thread, so the barrier is never initialized. /// /// # Example /// @@ -390,7 +414,8 @@ impl ManagedBarrier { /// ``` #[inline(always)] pub unsafe fn init_by(self, count: u32, init_thread: u32) -> ManagedBarrier { - if crate::thread::threadIdx_x() == init_thread { + // Block-linear id: x fastest, then y, then z; 0 is thread (0, 0, 0). + if crate::thread::block_linear_thread_id_here() == init_thread { unsafe { mbarrier_init(self.ptr as *mut Barrier, count); fence_proxy_async_shared_cta(); @@ -412,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. @@ -483,8 +535,9 @@ impl ManagedBarrier { /// Invalidate the barrier. /// - /// **All threads in the block should call this.** Only thread 0 performs - /// the actual invalidation; all threads synchronize before returning. + /// **All threads in the block should call this.** Only linear thread 0, + /// that is thread (0, 0, 0), performs the actual invalidation; all threads + /// synchronize before returning. /// /// This is a convenience wrapper for `inval_by(0)`. /// @@ -499,7 +552,7 @@ impl ManagedBarrier { /// # Example /// /// ```rust,ignore - /// // ALL threads call inval - only thread 0 actually invalidates + /// // ALL threads call inval - only linear thread 0 actually invalidates /// let _dead = unsafe { bar.inval() }; // Consumes Ready, returns Invalidated /// ``` #[inline(always)] @@ -509,21 +562,26 @@ impl ManagedBarrier { /// Invalidate the barrier with a specific thread performing invalidation. /// - /// **All threads in the block should call this.** Only the thread with - /// `threadIdx.x == inval_thread` performs the actual invalidation; + /// **All threads in the block should call this.** Only the thread whose + /// linear id equals `inval_thread` performs the actual invalidation; /// all threads synchronize before returning. /// /// Consumes the `Ready` barrier and returns an `Invalidated` barrier. /// /// # Parameters /// - /// - `inval_thread`: Thread ID (threadIdx.x) that performs invalidation + /// - `inval_thread`: Linear thread id within the block that performs + /// invalidation. `0` is thread `(0, 0, 0)`. In a 1D block this is + /// `threadIdx.x`. /// /// # Safety /// /// - All threads must have completed their wait operations before calling /// - All participating threads in the block must call this together - /// - `inval_thread` must be a valid thread ID within the block + /// - `inval_thread` must be the same value in every thread and below + /// `blockDim.x * blockDim.y * blockDim.z`. A selector that differs between + /// threads elects zero or several threads; an out-of-range id selects no + /// thread, so the barrier is never invalidated. /// /// # Example /// @@ -536,7 +594,8 @@ impl ManagedBarrier { // Ensure all threads are done with the barrier before invalidating crate::thread::sync_threads(); - if crate::thread::threadIdx_x() == inval_thread { + // Block-linear id: x fastest, then y, then z; 0 is thread (0, 0, 0). + if crate::thread::block_linear_thread_id_here() == inval_thread { unsafe { mbarrier_inval(self.ptr as *mut Barrier) }; } 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/cooperative_groups.rs b/cuda-oxide/crates/cuda-device/src/cooperative_groups.rs index 1c1005f086..07d50d592f 100644 --- a/cuda-oxide/crates/cuda-device/src/cooperative_groups.rs +++ b/cuda-oxide/crates/cuda-device/src/cooperative_groups.rs @@ -193,17 +193,6 @@ pub trait WarpCollective: ThreadGroup { // Helpers (not part of the public surface) // ============================================================================= -/// Linearize the (x, y, z) thread index inside the current block. -#[inline(always)] -fn thread_in_block_linear() -> u32 { - let tx = thread::threadIdx_x(); - let ty = thread::threadIdx_y(); - let tz = thread::threadIdx_z(); - let dx = thread::blockDim_x(); - let dy = thread::blockDim_y(); - (tz * dy + ty) * dx + tx -} - /// Total threads per block. #[inline(always)] fn threads_per_block() -> u32 { @@ -268,7 +257,7 @@ impl ThreadGroup for Grid { #[inline(always)] fn thread_rank(&self) -> u32 { - block_in_grid_linear() * threads_per_block() + thread_in_block_linear() + block_in_grid_linear() * threads_per_block() + thread::block_linear_thread_id_here() } #[inline(always)] @@ -319,7 +308,7 @@ impl ThreadGroup for Cluster { #[inline(always)] fn thread_rank(&self) -> u32 { - cluster::block_rank() * threads_per_block() + thread_in_block_linear() + cluster::block_rank() * threads_per_block() + thread::block_linear_thread_id_here() } #[inline(always)] @@ -383,7 +372,7 @@ impl ThreadGroup for ThreadBlock { #[inline(always)] fn thread_rank(&self) -> u32 { - thread_in_block_linear() + thread::block_linear_thread_id_here() } #[inline(always)] @@ -1279,13 +1268,6 @@ where acc } -/// Linear warp index inside the current block. Works for any block dim -/// shape; `warp::warp_id` only handles the 1D case. -#[inline(always)] -fn warp_in_block_linear() -> u32 { - thread_in_block_linear() / 32 -} - /// Reduce one value per thread down to a single value across a thread block. /// /// Three-phase shape (the same shape CUB and cuCollections use): @@ -1369,7 +1351,7 @@ where let smem = unsafe { SharedArray::as_raw_mut_ptr(smem) }; let lane = warp::lane_id(); - let warp_id = warp_in_block_linear(); + let warp_id = warp::block_linear_warp_id_here(); let remaining = block_threads - warp_id * 32; let live = if remaining < 32 { remaining } else { 32 }; @@ -1475,7 +1457,7 @@ where crate::gpu_assert!(runtime_warps as usize <= NUM_WARPS); let lane = warp::lane_id(); - let warp_id = warp_in_block_linear(); + let warp_id = warp::block_linear_warp_id_here(); let remaining = block_threads - warp_id * 32; let live = if remaining < 32 { remaining } else { 32 }; diff --git a/cuda-oxide/crates/cuda-device/src/debug.rs b/cuda-oxide/crates/cuda-device/src/debug.rs index 1d7c66869c..b300d2cdf4 100644 --- a/cuda-oxide/crates/cuda-device/src/debug.rs +++ b/cuda-oxide/crates/cuda-device/src/debug.rs @@ -194,14 +194,22 @@ pub fn __gpu_vprintf(_format: *const u8, _args: *const u8) -> i32 { unreachable!("__gpu_vprintf called outside CUDA kernel context") } -/// Trait for GPU printf argument promotion. +/// Conversion table for `gpu_printf!` arguments. /// -/// Implements C vararg promotion rules: +/// A placeholder without a type character (`{}`, `{:8}`, `{:08}`, `{:.2}`) +/// takes its C conversion from this trait: [`FORMAT_CHAR`](Self::FORMAT_CHAR) +/// is `d`, `u`, `f` or `p`, [`IS_64BIT`](Self::IS_64BIT) adds the `ll` length +/// modifier to an integer conversion, and [`IS_FLOAT`](Self::IS_FLOAT) selects +/// `%f`. [`promote`](Self::promote) applies the C vararg promotion rules before +/// the value is packed: /// - `i8`, `i16` → `i32` /// - `u8`, `u16` → `u32` /// - `f32` → `f64` /// - `bool` → `i32` -/// - 64-bit types stay as-is +/// - 64-bit integers and raw pointers stay as they are +/// +/// `char`, fieldless enums, `i128` and `u128` have no impl. Print a `char` +/// with `{:c}` and cast an enum with `as i32`. pub trait GpuPrintfArg { /// The promoted type for C vararg ABI type Promoted: Copy; @@ -374,5 +382,305 @@ impl GpuPrintfArg for *mut T { } } +// ============================================================================= +// Compile-time format string assembly +// ============================================================================= + +/// One conversion of a `gpu_printf!` format string. +/// +/// The macro hands a list of these to [`__gpu_printf_format`] together with +/// the literal text around them. Public only because the macro expands in +/// downstream crates. +#[doc(hidden)] +#[derive(Clone, Copy, Debug)] +pub enum __GpuPrintfSpec { + /// A complete C conversion, `%` included, chosen by the macro from an + /// explicit type character such as `{:x}` or `{:.2f}`. + Fixed(&'static [u8]), + /// A conversion whose type character comes from the argument's + /// [`GpuPrintfArg`] constants. `prefix` holds the flags, width and + /// precision that sit between the `%` and the type character. + Typed { + prefix: &'static [u8], + format_char: char, + is_64bit: bool, + is_float: bool, + }, +} + +/// Assembles the NUL-terminated C format string of one `gpu_printf!` call. +/// +/// `parts` is the literal text around the conversions, with `%` already +/// doubled, and `specs` the conversions in order, so `parts.len()` is +/// `specs.len() + 1`. A typed conversion gets `ll` for a 64-bit integer that +/// is not a pointer, and `f` for a floating-point argument. `N` is an upper +/// bound the macro computes from the same data; the bytes after the NUL stay +/// zero. +/// +/// The macro evaluates this at compile time, so each `assert!` and `panic!` +/// here is a compile error at the call site. Public only because the macro +/// expands in downstream crates. +#[doc(hidden)] +pub const fn __gpu_printf_format( + parts: &[&[u8]], + specs: &[__GpuPrintfSpec], +) -> [u8; N] { + assert!( + parts.len() == specs.len() + 1, + "gpu_printf: each conversion needs the literal text before and after it" + ); + + let mut out = [0u8; N]; + let mut at = 0; + let mut index = 0; + while index < parts.len() { + at = gpu_printf_copy(&mut out, at, parts[index]); + if index < specs.len() { + at = match specs[index] { + __GpuPrintfSpec::Fixed(bytes) => gpu_printf_copy(&mut out, at, bytes), + __GpuPrintfSpec::Typed { + prefix, + format_char, + is_64bit, + is_float, + } => gpu_printf_typed(&mut out, at, prefix, format_char, is_64bit, is_float), + }; + } + index += 1; + } + assert!( + at < N, + "gpu_printf: the format buffer is too small for its format string" + ); + out[at] = 0; + out +} + +/// Writes one byte of the format string. +const fn gpu_printf_put(out: &mut [u8], at: usize, byte: u8) -> usize { + assert!( + at < out.len(), + "gpu_printf: the format buffer is too small for its format string" + ); + out[at] = byte; + at + 1 +} + +/// Copies `bytes` into the format string. +const fn gpu_printf_copy(out: &mut [u8], mut at: usize, bytes: &[u8]) -> usize { + let mut index = 0; + while index < bytes.len() { + at = gpu_printf_put(out, at, bytes[index]); + index += 1; + } + at +} + +/// Writes a conversion whose type character comes from [`GpuPrintfArg`]. +/// +/// C defines the `+` and space flags for signed conversions only; on `%u` or +/// `%p` the output is undefined, so the combination is rejected here. +const fn gpu_printf_typed( + out: &mut [u8], + mut at: usize, + prefix: &[u8], + format_char: char, + is_64bit: bool, + is_float: bool, +) -> usize { + if !is_float && format_char != 'd' && gpu_printf_has_sign_flag(prefix) { + panic!( + "gpu_printf: the + and space flags need a signed integer or floating-point argument" + ); + } + at = gpu_printf_put(out, at, b'%'); + at = gpu_printf_copy(out, at, prefix); + if !is_float && is_64bit && format_char != 'p' { + at = gpu_printf_copy(out, at, b"ll"); + } + let type_char = if is_float { b'f' } else { format_char as u8 }; + gpu_printf_put(out, at, type_char) +} + +/// Whether a conversion prefix carries the `+` or the space flag. +const fn gpu_printf_has_sign_flag(prefix: &[u8]) -> bool { + let mut index = 0; + while index < prefix.len() { + if prefix[index] == b'+' || prefix[index] == b' ' { + return true; + } + index += 1; + } + false +} + // Re-export the gpu_printf macro from cuda-macros pub use cuda_macros::gpu_printf; + +#[cfg(test)] +mod tests { + use super::{__GpuPrintfSpec, __gpu_printf_format, gpu_printf}; + + fn typed( + prefix: &'static [u8], + format_char: char, + is_64bit: bool, + is_float: bool, + ) -> __GpuPrintfSpec { + __GpuPrintfSpec::Typed { + prefix, + format_char, + is_64bit, + is_float, + } + } + + /// The bytes before the NUL, as text. + fn text(buffer: &[u8]) -> &str { + let end = buffer + .iter() + .position(|&byte| byte == 0) + .expect("format strings are NUL-terminated"); + core::str::from_utf8(&buffer[..end]).expect("format strings are ASCII") + } + + #[test] + fn typed_conversions_follow_the_trait_constants() { + let cases: [(&[u8], char, bool, bool, &str); 8] = [ + (b"", 'd', false, false, "%d"), + (b"", 'd', true, false, "%lld"), + (b"", 'u', false, false, "%u"), + (b"", 'u', true, false, "%llu"), + (b"", 'f', false, true, "%f"), + (b"", 'p', true, false, "%p"), + (b".2", 'f', false, true, "%.2f"), + (b".3", 'd', false, false, "%.3d"), + ]; + for (prefix, format_char, is_64bit, is_float, expected) in cases { + let buffer = __gpu_printf_format::<8>( + &[b"", b""], + &[typed(prefix, format_char, is_64bit, is_float)], + ); + assert_eq!(text(&buffer), expected); + } + } + + #[test] + fn prefixes_sit_between_the_percent_and_the_type_character() { + let buffer = __gpu_printf_format::<32>( + &[b"[", b"] [", b"] [", b"]\n"], + &[ + typed(b"-8", 'd', false, false), + typed(b"08", 'u', true, false), + typed(b"*", 'd', true, false), + ], + ); + assert_eq!(text(&buffer), "[%-8d] [%08llu] [%*lld]\n"); + } + + #[test] + fn fixed_conversions_and_literal_text_are_copied_verbatim() { + let buffer = __gpu_printf_format::<32>( + &[b"100%% done: ", b" ", b"\n"], + &[ + __GpuPrintfSpec::Fixed(b"%#llx"), + typed(b"", 'f', false, true), + ], + ); + assert_eq!(text(&buffer), "100%% done: %#llx %f\n"); + } + + #[test] + fn the_buffer_may_be_filled_to_the_last_byte() { + let buffer = __gpu_printf_format::<5>(&[b"", b""], &[typed(b"", 'd', true, false)]); + assert_eq!(buffer, *b"%lld\0"); + } + + #[test] + fn bytes_after_the_nul_stay_zero() { + // The macro leaves room for `ll`; a 32-bit argument does not use it. + let buffer = __gpu_printf_format::<5>(&[b"", b""], &[typed(b"", 'd', false, false)]); + assert_eq!(buffer, *b"%d\0\0\0"); + } + + #[test] + fn builds_in_a_const_context() { + const FORMAT: [u8; 8] = __gpu_printf_format::<8>( + &[b"v=", b"\n"], + &[__GpuPrintfSpec::Typed { + prefix: b"", + format_char: 'u', + is_64bit: true, + is_float: false, + }], + ); + assert_eq!(text(&FORMAT), "v=%llu\n"); + } + + #[test] + fn sign_flags_are_kept_on_signed_and_float_arguments() { + let buffer = __gpu_printf_format::<16>( + &[b"", b" ", b" ", b""], + &[ + typed(b"+", 'd', true, false), + typed(b" ", 'd', false, false), + typed(b"+", 'f', false, true), + ], + ); + assert_eq!(text(&buffer), "%+lld % d %+f"); + } + + #[test] + #[should_panic(expected = "the + and space flags need a signed integer or floating-point")] + fn the_plus_flag_on_an_unsigned_argument_is_rejected() { + let _ = __gpu_printf_format::<8>(&[b"", b""], &[typed(b"+", 'u', false, false)]); + } + + #[test] + #[should_panic(expected = "the + and space flags need a signed integer or floating-point")] + fn the_space_flag_on_a_pointer_argument_is_rejected() { + let _ = __gpu_printf_format::<8>(&[b"", b""], &[typed(b" ", 'p', true, false)]); + } + + #[test] + #[should_panic(expected = "the format buffer is too small")] + fn a_buffer_too_small_for_the_text_is_rejected() { + let _ = __gpu_printf_format::<4>(&[b"", b""], &[typed(b"", 'd', true, false)]); + } + + #[test] + #[should_panic(expected = "the format buffer is too small")] + fn a_buffer_with_no_room_for_the_nul_is_rejected() { + let _ = __gpu_printf_format::<2>(&[b"", b""], &[typed(b"", 'd', false, false)]); + } + + #[test] + #[should_panic(expected = "each conversion needs the literal text before and after it")] + fn parts_and_conversions_must_interleave() { + let _ = __gpu_printf_format::<8>(&[b""], &[typed(b"", 'd', false, false)]); + } + + /// This call is reached, so the expansion is monomorphized and the inline + /// const that builds its format string is evaluated while the test binary + /// compiles. A size mismatch between the macro and the builder, or a + /// missing `GpuPrintfArg` impl, fails this test at build time. + /// `__gpu_vprintf` then panics because there is no kernel around it. + #[test] + #[should_panic(expected = "__gpu_vprintf called outside CUDA kernel context")] + fn default_placeholders_expand_and_build_their_format_on_the_host() { + let value: u64 = u64::MAX; + let ratio: f32 = 1.25; + let pointer: *const f32 = ∶ + gpu_printf!( + "{} {:.2} {:08} {} {} {:x} 100% [{:*}]\n", + value, + ratio, + 42, + true, + pointer, + 255u32, + 6, + -7i64 + ); + } +} 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 e44c9a333c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 47a5053d24..54e963cc62 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/atomic.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/atomic.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fb8215c889..85254d74f7 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/bf16.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/bf16.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 b9f4be0b75..2f952e598e 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/bf16x2.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/bf16x2.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 94c9cae176..d11649381e 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/clc.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/clc.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 9115aec1b5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 d665c6fc59..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 4e2ce888c0..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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`. // Its private leaves let public helpers compose generated reads. diff --git a/cuda-oxide/crates/cuda-device/src/generated/convert.rs b/cuda-oxide/crates/cuda-device/src/generated/convert.rs index f84862afdd..f592b2299e 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/convert.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/convert.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. // Included inside `cuda_device::convert`. @@ -81,6 +81,15 @@ pub fn cvt_rn_relu_f16x2_f32(lo: f32, hi: f32) -> u32 { ) } +/// Converts packed bf16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order. +#[inline(never)] +pub fn cvt_rn_satfinite_e4m3x2_bf16x2(packed: u32) -> u16 { + let _ = packed; + unreachable!( + "generated CUDA intrinsic `cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2` executed outside device compilation" + ) +} + /// Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order. #[inline(never)] pub fn cvt_rn_satfinite_e4m3x2_f16x2(packed: u32) -> u16 { @@ -99,6 +108,15 @@ pub fn cvt_rn_satfinite_e4m3x2_f32(lo: f32, hi: f32) -> u16 { ) } +/// Converts packed bf16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order. +#[inline(never)] +pub fn cvt_rn_satfinite_e5m2x2_bf16x2(packed: u32) -> u16 { + let _ = packed; + unreachable!( + "generated CUDA intrinsic `cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2` executed outside device compilation" + ) +} + /// Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order. #[inline(never)] pub fn cvt_rn_satfinite_e5m2x2_f16x2(packed: u32) -> u16 { @@ -117,6 +135,15 @@ pub fn cvt_rn_satfinite_e5m2x2_f32(lo: f32, hi: f32) -> u16 { ) } +/// Converts packed bf16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order. +#[inline(never)] +pub fn cvt_rn_satfinite_relu_e4m3x2_bf16x2(packed: u32) -> u16 { + let _ = packed; + unreachable!( + "generated CUDA intrinsic `cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2` executed outside device compilation" + ) +} + /// Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order. #[inline(never)] pub fn cvt_rn_satfinite_relu_e4m3x2_f16x2(packed: u32) -> u16 { @@ -135,6 +162,15 @@ pub fn cvt_rn_satfinite_relu_e4m3x2_f32(lo: f32, hi: f32) -> u16 { ) } +/// Converts packed bf16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order. +#[inline(never)] +pub fn cvt_rn_satfinite_relu_e5m2x2_bf16x2(packed: u32) -> u16 { + let _ = packed; + unreachable!( + "generated CUDA intrinsic `cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2` executed outside device compilation" + ) +} + /// Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order. #[inline(never)] pub fn cvt_rn_satfinite_relu_e5m2x2_f16x2(packed: u32) -> u16 { 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 877e00f15f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6f52412e9a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 2028a15d31..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 2e1e16f829..d39b60563a 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/dotprod.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/dotprod.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 e0308d047a..acdeaf41b5 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/f16.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/f16.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 00085fa0b6..fa908ce6b5 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/f16x2.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/f16x2.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c5de9a8539..6769a157a2 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/f32x2.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/f32x2.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 685f7e4c38..06ef909025 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/fence.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/fence.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fbe53bd878..3e97a682d8 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/float.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/float.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 94d7848ee7..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 ed3fa06aeb..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 054da50cd7..77feb7b3eb 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/i16x2.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/i16x2.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 87a798b163..d13952b077 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/int.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/int.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 b288951019..1645ec9922 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/ldmatrix.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/ldmatrix.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 eb377051e6..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 7787d693cf..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: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 09ccf2ea4d..90c6bdff83 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/movmatrix.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/movmatrix.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 30fb07902d..2c3ce2bedf 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/prmt.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/prmt.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 91d3e7f239..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 f623acbe07..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -43,7 +43,7 @@ pub unsafe fn mma_m16n8k128_s32_b1_xor_popc(c: [i32; 4], a: [u32; 2], b: u32) -> unreachable!("mma_m16n8k128_s32_b1_xor_popc called outside CUDA kernel context") } -/// Executes the mma_m16n8k16_f16_f16 Ampere floating-point MMA form. +/// Executes the mma_m16n8k16_f16_f16 dense floating-point MMA form. /// Lowers to `mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16`. C, A, and B are this lane's fragments in PTX register order. /// Requires PTX 7.0 and `sm_80+`. /// @@ -2515,7 +2515,7 @@ pub unsafe fn mma_m16n8k32_s32_u8_satfinite(c: [i32; 4], a: [u32; 4], b: [u32; 2 unreachable!("mma_m16n8k32_s32_u8_satfinite called outside CUDA kernel context") } -/// Executes the mma_m16n8k4_f32_tf32 Ampere floating-point MMA form. +/// Executes the mma_m16n8k4_f32_tf32 dense floating-point MMA form. /// Lowers to `mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32`. C, A, and B are this lane's fragments in PTX register order. /// Requires PTX 7.0 and `sm_80+`. /// @@ -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+`. @@ -2667,9 +2706,9 @@ pub unsafe fn mma_m16n8k64_s32_u4_satfinite(c: [i32; 4], a: [u32; 4], b: [u32; 2 unreachable!("mma_m16n8k64_s32_u4_satfinite called outside CUDA kernel context") } -/// Executes the mma_m16n8k8_f16_f16 Ampere floating-point MMA form. +/// Executes the mma_m16n8k8_f16_f16 dense floating-point MMA form. /// Lowers to `mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16`. C, A, and B are this lane's fragments in PTX register order. -/// Requires PTX 7.0 and `sm_80+`. +/// Requires PTX 6.5 and `sm_75+`. /// /// # Safety /// All 32 lanes must execute the same instruction with the same qualifiers, and no lane may have exited. @@ -2683,7 +2722,7 @@ pub unsafe fn mma_m16n8k8_f16_f16(c: [u32; 2], a: [u32; 2], b: u32) -> [u32; 2] unreachable!("mma_m16n8k8_f16_f16 called outside CUDA kernel context") } -/// Executes the mma_m16n8k8_f32_bf16 Ampere floating-point MMA form. +/// Executes the mma_m16n8k8_f32_bf16 dense floating-point MMA form. /// Lowers to `mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32`. C, A, and B are this lane's fragments in PTX register order. /// Requires PTX 7.0 and `sm_80+`. /// @@ -2699,9 +2738,9 @@ pub unsafe fn mma_m16n8k8_f32_bf16(c: [f32; 4], a: [u32; 2], b: u32) -> [f32; 4] unreachable!("mma_m16n8k8_f32_bf16 called outside CUDA kernel context") } -/// Executes the mma_m16n8k8_f32_f16 Ampere floating-point MMA form. +/// Executes the mma_m16n8k8_f32_f16 dense floating-point MMA form. /// Lowers to `mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32`. C, A, and B are this lane's fragments in PTX register order. -/// Requires PTX 7.0 and `sm_80+`. +/// Requires PTX 6.5 and `sm_75+`. /// /// # Safety /// All 32 lanes must execute the same instruction with the same qualifiers, and no lane may have exited. 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 35198aaa6f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 a1e8bd0f63..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -416,7 +416,7 @@ pub unsafe fn mma_sp_m16n8k32_s32_u8_satfinite( /// # 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 4-, 4-, and 4-register fragments; `metadata` must contain its sparse metadata. -/// Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior. +/// Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior. /// `selector` must be the compile-time constant `0`. /// See the [PTX sparse MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp). /// This register-only operation is not a memory fence. @@ -440,7 +440,7 @@ pub unsafe fn mma_sp_m16n8k64_f32_e4m3_e4m3_f32( /// # 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 4-, 4-, and 4-register fragments; `metadata` must contain its sparse metadata. -/// Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior. +/// Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior. /// `selector` must be the compile-time constant `0`. /// See the [PTX sparse MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp). /// This register-only operation is not a memory fence. @@ -464,7 +464,7 @@ pub unsafe fn mma_sp_m16n8k64_f32_e4m3_e5m2_f32( /// # 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 4-, 4-, and 4-register fragments; `metadata` must contain its sparse metadata. -/// Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior. +/// Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior. /// `selector` must be the compile-time constant `0`. /// See the [PTX sparse MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp). /// This register-only operation is not a memory fence. @@ -488,7 +488,7 @@ pub unsafe fn mma_sp_m16n8k64_f32_e5m2_e4m3_f32( /// # 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 4-, 4-, and 4-register fragments; `metadata` must contain its sparse metadata. -/// Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior. +/// Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior. /// `selector` must be the compile-time constant `0`. /// See the [PTX sparse MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp). /// This register-only operation is not a memory fence. diff --git a/cuda-oxide/crates/cuda-device/src/generated/sreg.rs b/cuda-oxide/crates/cuda-device/src/generated/sreg.rs index f972a0a046..d279384792 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/sreg.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/sreg.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 // DefPaths remain stable during the generated-intrinsics migration. diff --git a/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs b/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs index 73f22a70d0..cba095423c 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/stmatrix.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 cb1712552e..d2ae333e1a 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/tcgen05.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/tcgen05.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 36156f0b8d..40b2eae7d5 100644 --- a/cuda-oxide/crates/cuda-device/src/generated/tma.rs +++ b/cuda-oxide/crates/cuda-device/src/generated/tma.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -233,6 +233,23 @@ pub unsafe fn cp_async_bulk_tensor_1d_g2s( unreachable!("cp_async_bulk_tensor_1d_g2s called outside CUDA kernel context") } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// # Safety +/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. +/// `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +#[inline(never)] +pub unsafe fn cp_async_bulk_tensor_1d_g2s_cta( + dst: SharedPtr, + tensor_map: *const TmaDescriptor, + coord0: i32, + barrier: *mut Barrier, +) { + let _ = (dst, tensor_map, coord0, barrier); + unreachable!("cp_async_bulk_tensor_1d_g2s_cta called outside CUDA kernel context") +} + /// Starts a TMA tile copy from shared to global memory. /// /// # Safety @@ -263,6 +280,24 @@ pub unsafe fn cp_async_bulk_tensor_2d_g2s( unreachable!("cp_async_bulk_tensor_2d_g2s called outside CUDA kernel context") } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// # Safety +/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. +/// `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +#[inline(never)] +pub unsafe fn cp_async_bulk_tensor_2d_g2s_cta( + dst: SharedPtr, + tensor_map: *const TmaDescriptor, + coord0: i32, + coord1: i32, + barrier: *mut Barrier, +) { + let _ = (dst, tensor_map, coord0, coord1, barrier); + unreachable!("cp_async_bulk_tensor_2d_g2s_cta called outside CUDA kernel context") +} + /// Starts a multicast TMA tile copy from global to cluster shared memory. /// /// # Safety @@ -329,6 +364,25 @@ pub unsafe fn cp_async_bulk_tensor_3d_g2s( unreachable!("cp_async_bulk_tensor_3d_g2s called outside CUDA kernel context") } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// # Safety +/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. +/// `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +#[inline(never)] +pub unsafe fn cp_async_bulk_tensor_3d_g2s_cta( + dst: SharedPtr, + tensor_map: *const TmaDescriptor, + coord0: i32, + coord1: i32, + coord2: i32, + barrier: *mut Barrier, +) { + let _ = (dst, tensor_map, coord0, coord1, coord2, barrier); + unreachable!("cp_async_bulk_tensor_3d_g2s_cta called outside CUDA kernel context") +} + /// Starts a TMA tile copy from shared to global memory. /// /// # Safety @@ -363,6 +417,26 @@ pub unsafe fn cp_async_bulk_tensor_4d_g2s( unreachable!("cp_async_bulk_tensor_4d_g2s called outside CUDA kernel context") } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// # Safety +/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. +/// `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +#[inline(never)] +pub unsafe fn cp_async_bulk_tensor_4d_g2s_cta( + dst: SharedPtr, + tensor_map: *const TmaDescriptor, + coord0: i32, + coord1: i32, + coord2: i32, + coord3: i32, + barrier: *mut Barrier, +) { + let _ = (dst, tensor_map, coord0, coord1, coord2, coord3, barrier); + unreachable!("cp_async_bulk_tensor_4d_g2s_cta called outside CUDA kernel context") +} + /// Starts a TMA tile copy from shared to global memory. /// /// # Safety @@ -402,6 +476,30 @@ pub unsafe fn cp_async_bulk_tensor_5d_g2s( unreachable!("cp_async_bulk_tensor_5d_g2s called outside CUDA kernel context") } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// # Safety +/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes. +/// `barrier` must point into the issuing CTA's shared memory, as for the other +/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type. +#[inline(never)] +#[allow(clippy::too_many_arguments)] +pub unsafe fn cp_async_bulk_tensor_5d_g2s_cta( + dst: SharedPtr, + tensor_map: *const TmaDescriptor, + coord0: i32, + coord1: i32, + coord2: i32, + coord3: i32, + coord4: i32, + barrier: *mut Barrier, +) { + let _ = ( + dst, tensor_map, coord0, coord1, coord2, coord3, coord4, barrier, + ); + unreachable!("cp_async_bulk_tensor_5d_g2s_cta called outside CUDA kernel context") +} + /// Starts a TMA tile copy from shared to global memory. /// /// # Safety 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 8019adef7b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fd1ea8426e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 a3171afbcf..67ee46790e 100644 --- a/cuda-oxide/crates/cuda-device/src/lib.rs +++ b/cuda-oxide/crates/cuda-device/src/lib.rs @@ -12,7 +12,7 @@ extern crate self as cuda_device; pub use cuda_macros::{ cluster_launch, constant, convergent, cooperative_launch, cuda_module, device, gpu_printf, - kernel, launch_bounds, launch_contract, ptx_asm, pure, readonly, + kernel, launch_bounds, launch_contract, ptx_asm, pure, readonly, tile_export, }; // Re-export for convenience @@ -82,7 +82,7 @@ pub use cusimd::{CuSimd, Float2, Float4, TmemRegs4, TmemRegs32}; pub use disjoint::{__LaunchContractDisjointSlice, __LaunchContractDisjointSliceAbi}; pub use disjoint::{DisjointSlice, SpaceLayout}; pub use fence::*; -pub use shared::{DynamicSharedArray, SharedArray}; +pub use shared::{DynamicSharedArray, SharedArray, SharedPtr}; pub use tcgen05::{ TensorMemoryHandle, TmemAddress, TmemDeallocated, TmemF32x4, TmemF32x32, TmemGuard, TmemReady, TmemUninit, diff --git a/cuda-oxide/crates/cuda-device/src/shared.rs b/cuda-oxide/crates/cuda-device/src/shared.rs index 5467ee7d30..a6a7ed0c66 100644 --- a/cuda-oxide/crates/cuda-device/src/shared.rs +++ b/cuda-oxide/crates/cuda-device/src/shared.rs @@ -80,7 +80,8 @@ use core::ops::{Index, IndexMut}; /// /// - `T`: Element type (f32, f64, i32, etc.) /// - `N`: Array size (must be known at compile time) -/// - `ALIGN`: Alignment in bytes (0 = natural alignment, 128 for TMA destinations) +/// - `ALIGN`: Alignment in bytes. `0` means natural alignment; any other value +/// must be a power of two (128 for TMA destinations). /// /// # Alignment /// @@ -96,6 +97,18 @@ use core::ops::{Index, IndexMut}; /// static mut TMA_TILE: SharedArray = SharedArray::UNINIT; /// ``` /// +/// The contract the compiler enforces: +/// +/// - `ALIGN` is `0` or a non-zero power of two. Anything else is a compile +/// error that names the use site. +/// - Storage is never declared below `align_of::()`. A request below it is +/// raised silently, as `#[repr(align(N))]` raises a struct; a larger request +/// is honored. So `SharedArray` is declared at 8 bytes, and a +/// `#[repr(C, align(16))]` element at natural alignment is declared at 16. +/// - Aggregate, nested-array and `f16` elements at natural alignment declare +/// their natural alignment in PTX (`.align 4` for `[f32; N]`, `.align 2` for +/// `f16`). Older releases printed a blanket `.align 8` for them. +/// /// # Soundness: `!Sync` /// /// `SharedArray` is intentionally `!Sync`. GPU shared memory is concurrently @@ -220,6 +233,34 @@ impl SharedArray { let _ = shared; unreachable!("SharedArray::as_raw_mut_ptr called outside CUDA kernel context") } + + /// Returns a [`SharedPtr`] to the first element of a shared array. + /// + /// Unlike [`Self::as_raw_mut_ptr`], which returns a generic pointer, the + /// result keeps the array's CTA-shared address space in its type, so it + /// stays a shared-memory pointer when it is passed to helper functions. + /// + /// # Safety + /// + /// `shared` must point to a `static mut SharedArray` in the current CUDA + /// kernel, normally obtained with `&raw mut`. The returned pointer is valid + /// only within that kernel. This requirement is what makes the type's + /// claim true: a `SharedPtr` built from any other pointer hands + /// [`SharedPtr::addr_u32`] an address outside the shared window, where the + /// conversion yields an unspecified value. Callers must keep concurrent + /// accesses disjoint and use the CUDA synchronization required before + /// reading data written by another thread. + /// + /// ```rust,ignore + /// static mut TILE: SharedArray = SharedArray::UNINIT; + /// let tile: SharedPtr = unsafe { SharedArray::shared_ptr(&raw mut TILE) }; + /// ``` + #[inline(always)] + pub unsafe fn shared_ptr(shared: *mut Self) -> SharedPtr { + SharedPtr { + ptr: shared as *mut SharedSlot, + } + } } impl Index for SharedArray { @@ -305,7 +346,8 @@ pub fn __dynamic_shared_alignment() { /// # Type Parameters /// /// - `T`: Element type for the returned pointer -/// - `ALIGN`: Base alignment in bytes (default: 16, matching nvcc) +/// - `ALIGN`: Base alignment in bytes (default: 16, matching nvcc). Must be a +/// non-zero power of two. /// /// # Alignment /// @@ -331,6 +373,18 @@ pub fn __dynamic_shared_alignment() { /// - `128`: Required for TMA (Tensor Memory Accelerator) operations /// - `256`: Cache-line friendly for some architectures /// +/// The contract the compiler enforces: +/// +/// - `ALIGN` is a non-zero power of two. Anything else, including `0`, is a +/// compile error that names the use site. +/// - A kernel has one dynamic pool. It is declared at the strongest typed +/// requirement over the kernel's `get`, `get_raw` and `offset` uses, each +/// taken as the larger of its `ALIGN` and `align_of::()`, combined with +/// the launch contract's alignment. `DynamicSharedArray::::get()` +/// alone declares the pool at 8 bytes. +/// - The declaration aligns the pool base. Keeping an `offset(byte_offset)` +/// result aligned for `T` is still the caller's responsibility. +/// /// # Usage /// /// ```rust,ignore @@ -484,8 +538,9 @@ impl DynamicSharedArray { /// Ensure `byte_offset` is properly aligned for type `T`. For example, /// `f32` requires 4-byte alignment, `f64` requires 8-byte alignment. /// - /// The base of dynamic shared memory is aligned to `ALIGN` bytes (default 16), - /// so to get an N-byte aligned pointer, ensure `byte_offset % N == 0`. + /// The base of dynamic shared memory is aligned to at least `ALIGN` bytes + /// (default 16) and never below `align_of::()`, so to get an N-byte + /// aligned pointer, ensure `byte_offset % N == 0`. #[inline(never)] pub fn offset(byte_offset: usize) -> *mut T { // The byte_offset is used by the compiler to compute the pointer. @@ -493,6 +548,24 @@ impl DynamicSharedArray { let _ = byte_offset; unreachable!("DynamicSharedArray::offset called outside CUDA kernel context") } + + /// Like [`Self::get`], but returns a [`SharedPtr`], which keeps the + /// CTA-shared address space in its type. + #[inline(always)] + pub fn shared_ptr() -> SharedPtr { + SharedPtr { + ptr: Self::get() as *mut SharedSlot, + } + } + + /// Like [`Self::offset`], but returns a [`SharedPtr`], which keeps the + /// CTA-shared address space in its type. + #[inline(always)] + pub fn shared_ptr_at(byte_offset: usize) -> SharedPtr { + SharedPtr { + ptr: Self::offset(byte_offset) as *mut SharedSlot, + } + } } /// Convert a generic-address pointer into its raw `.shared` window offset. @@ -526,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 @@ -546,4 +633,136 @@ pub unsafe fn cvta_generic_to_shared_u32(ptr: *const u8) -> u32 { unreachable!("cvta_generic_to_shared_u32 called outside CUDA kernel context") } +// ============================================================================ +// SharedPtr - a pointer into the issuing CTA's shared memory, typed as such +// ============================================================================ + +/// The pointee of [`SharedPtr`]'s field. A pointer to `SharedSlot` is a +/// pointer into CTA-shared memory (`addrspace(3)`) by its type, as a pointer +/// to [`SharedArray`] is; `repr(transparent)` gives it `T`'s size, so pointer +/// arithmetic strides by `T`. +#[doc(hidden)] +#[repr(transparent)] +pub struct SharedSlot(T); + +/// A raw pointer into the shared memory of the issuing CTA (`.shared::cta`). +/// +/// A `*mut T` in device code is a generic pointer, which may address global, +/// local or shared memory, including another CTA's shared memory in the +/// cluster. `SharedPtr` instead carries its address space in its type: the +/// compiler knows it is a CTA-shared pointer wherever it goes, including +/// through calls to helper functions, and emits the `.shared::cta` forms of +/// instructions that take it without converting a generic address. +/// +/// It is created from shared memory the CTA owns, with the `unsafe` +/// [`SharedArray::shared_ptr`] or the safe [`DynamicSharedArray::shared_ptr`], +/// or from a generic pointer with the `unsafe` [`SharedPtr::from_raw`]. Like a raw +/// pointer, it may dangle; dereferencing it goes through [`SharedPtr::as_ptr`] +/// and `unsafe` code. +/// +/// Pointers from [`crate::cluster::map_shared_rank`], which may name another +/// CTA's shared memory, are not `SharedPtr`s. +/// +/// # Example +/// +/// ```rust,ignore +/// static mut TILE: SharedArray = SharedArray::UNINIT; +/// +/// fn store(tile: SharedPtr, i: usize, value: f32) { +/// // `tile` is still a shared pointer here, although `store` is a +/// // separate function. +/// unsafe { tile.add(i).as_ptr().write(value) }; +/// } +/// +/// let tile = unsafe { SharedArray::shared_ptr(&raw mut TILE) }; +/// store(tile, thread::threadIdx_x() as usize, 1.0); +/// ``` +#[repr(transparent)] +pub struct SharedPtr { + ptr: *mut SharedSlot, +} + +impl Clone for SharedPtr { + #[inline(always)] + fn clone(&self) -> Self { + *self + } +} + +impl Copy for SharedPtr {} + +impl SharedPtr { + /// Converts a generic pointer to a `SharedPtr`, as CUDA C++'s + /// `__cvta_generic_to_shared` or CuTe's `make_smem_ptr` do. + /// + /// # Safety + /// + /// `ptr` must point into the shared memory of the issuing CTA. A pointer + /// to global or local memory, or to another CTA's shared memory, gives an + /// unspecified address. + #[inline(always)] + pub unsafe fn from_raw(ptr: *mut T) -> Self { + Self { + ptr: ptr as *mut SharedSlot, + } + } + + /// The generic pointer to the same memory, for ordinary loads and stores + /// and for APIs that take raw pointers. + #[inline(always)] + pub fn as_ptr(self) -> *mut T { + self.ptr as *mut T + } + + /// The pointer's 32-bit `.shared::cta` address, as taken by the + /// `_shared_u32` forms of CTA-shared instructions. + #[inline(always)] + pub fn addr_u32(self) -> u32 { + // SAFETY: `self` points into the issuing CTA's shared memory. + unsafe { cvta_generic_to_shared_u32(self.as_ptr() as *const u8) } + } + + /// Reinterprets the pointee type, keeping the address space. + #[inline(always)] + pub fn cast(self) -> SharedPtr { + SharedPtr { + ptr: self.ptr as *mut SharedSlot, + } + } + + /// The pointer `count` elements of `T` further on, as `<*mut T>::add`. + /// + /// # Safety + /// + /// As for `<*mut T>::add`. + #[inline(always)] + pub unsafe fn add(self, count: usize) -> Self { + Self { + ptr: unsafe { self.ptr.add(count) }, + } + } + + /// The pointer `count` elements of `T` away, as `<*mut T>::offset`. + /// + /// # Safety + /// + /// As for `<*mut T>::offset`. + #[inline(always)] + pub unsafe fn offset(self, count: isize) -> Self { + Self { + ptr: unsafe { self.ptr.offset(count) }, + } + } + + /// The pointer `count` bytes further on, as `<*mut T>::byte_add`. + /// + /// # Safety + /// + /// As for `<*mut T>::byte_add`. + #[inline(always)] + pub unsafe fn byte_add(self, count: usize) -> Self { + unsafe { self.cast::().add(count).cast::() } + } +} + include!("generated/shared_sreg.rs"); diff --git a/cuda-oxide/crates/cuda-device/src/tcgen05.rs b/cuda-oxide/crates/cuda-device/src/tcgen05.rs index 5970a81bf5..4254fb4551 100644 --- a/cuda-oxide/crates/cuda-device/src/tcgen05.rs +++ b/cuda-oxide/crates/cuda-device/src/tcgen05.rs @@ -46,7 +46,10 @@ //! use cuda_device::tcgen05::*; //! use cuda_device::barrier::*; //! -//! // 1. Allocate TMEM (warp-synchronous - all 32 threads in warp) +//! // 1. Allocate TMEM (warp-synchronous - all 32 threads in warp). +//! // `warp_id()` is `threadIdx.x / 32` and names one hardware warp only in +//! // a 1D block. A 2D or 3D block must select the warp by block-linear +//! // thread id, which `TmemGuard::alloc` / `dealloc` do. //! if warp_id() == 0 { //! tcgen05_alloc(tmem_addr_ptr, 64); // 64 columns //! } @@ -62,7 +65,7 @@ //! // 3. All threads wait via mbarrier //! mbarrier_try_wait(&mbar, 0); //! -//! // 4. Deallocate TMEM (warp-synchronous) +//! // 4. Deallocate TMEM (warp-synchronous; the 1D-only note from step 1 applies) //! if warp_id() == 0 { //! tcgen05_dealloc(tmem_addr, 64); //! } @@ -260,13 +263,23 @@ impl TmemGuard { /// /// # Parameters /// - /// - `alloc_warp`: Warp ID that performs allocation (all 32 threads in this warp) + /// - `alloc_warp`: Block-linear warp id that performs allocation (all 32 + /// threads in that warp). The id is `linear_tid / 32`, where + /// `linear_tid = threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z)`. + /// Warp 0 is linear tids `0..=31`. In a 1D block that matches + /// `threadIdx.x / 32`; in a 2D or 3D block it is not every row that + /// shares the same `threadIdx.x`. A kernel that allocates more than + /// once must issue every allocation from the same warp; the CUTLASS + /// TMEM allocator states this as a precondition. /// /// # Safety /// /// - Must be called before any TMEM operations /// - All participating threads in the block must call this together - /// - `alloc_warp` must be a valid warp ID within the block + /// - `alloc_warp` must name a full 32-thread warp: the block must hold at + /// least `32 * (alloc_warp + 1)` threads. The instruction is + /// `.sync.aligned`, and its behavior is undefined when a lane of the + /// warp has exited (PTX ISA 9.7.18.7.1, `tcgen05.alloc`) /// /// # Example /// @@ -276,7 +289,7 @@ impl TmemGuard { /// ``` #[inline(always)] pub unsafe fn alloc_by(self, alloc_warp: u32) -> TmemGuard { - let warp_id = crate::warp::warp_id(); + let warp_id = crate::warp::block_linear_warp_id_here(); if warp_id == alloc_warp { unsafe { tcgen05_alloc(self.smem_ptr, N_COLS) }; @@ -318,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 @@ -357,18 +372,28 @@ impl TmemGuard { /// /// # Parameters /// - /// - `dealloc_warp`: Warp ID that performs deallocation (must match alloc warp) + /// - `dealloc_warp`: Block-linear warp id that performs deallocation + /// (`linear_tid / 32`, same numbering as [`TmemGuard::alloc_by`]). Any + /// one full warp of the CTA may deallocate once every TMEM access is + /// complete; using the allocating warp is the convention here. The PTX + /// ISA asks for one warp of the CTA to perform the allocation and one + /// to perform the deallocation; it does not require them to be the + /// same warp. CUTLASS requires the same warp only for repeated + /// allocations within one kernel. /// /// # Safety /// /// - All threads must have completed their TMEM operations before calling /// - All participating threads in the block must call this together - /// - `dealloc_warp` should typically match the warp that allocated + /// - `dealloc_warp` must name a full 32-thread warp: the block must hold + /// at least `32 * (dealloc_warp + 1)` threads. The instruction is + /// `.sync.aligned`, and its behavior is undefined when a lane of the + /// warp has exited (PTX ISA 9.7.18.7.1, `tcgen05.dealloc`) /// /// # Example /// /// ```rust,ignore - /// // Use warp 1 for deallocation (should match allocation warp) + /// // Use warp 1 for deallocation; the allocating warp is the convention /// let _dead = unsafe { tmem.dealloc_by(1) }; /// ``` #[inline(always)] @@ -376,7 +401,7 @@ impl TmemGuard { // Ensure all threads are done with TMEM before deallocating crate::thread::sync_threads(); - let warp_id = crate::warp::warp_id(); + let warp_id = crate::warp::block_linear_warp_id_here(); if warp_id == dealloc_warp { unsafe { diff --git a/cuda-oxide/crates/cuda-device/src/thread.rs b/cuda-oxide/crates/cuda-device/src/thread.rs index 8b65b380c7..892ecf8d5f 100644 --- a/cuda-oxide/crates/cuda-device/src/thread.rs +++ b/cuda-oxide/crates/cuda-device/src/thread.rs @@ -1067,6 +1067,52 @@ pub fn index_2d_col() -> usize { include!("generated/sreg.rs"); include!("generated/register_control.rs"); +// ============================================================================= +// Block-Linear Thread Id +// ============================================================================= + +/// Position of a thread inside its block when the threads are numbered the +/// way the hardware forms warps: `x` fastest, then `y`, then `z`. +/// +/// ```text +/// linear_tid = thread_idx_x + block_dim_x * (thread_idx_y + block_dim_y * thread_idx_z) +/// ``` +/// +/// Each hardware warp holds 32 consecutive linear ids, so `linear_tid / 32` +/// names the warp a thread belongs to in any block layout. Works for any +/// block dim; [`crate::warp::warp_id`] is `threadIdx.x / 32` and only handles +/// the 1D case. `blockDim.z` is not an input because `z` is the slowest axis. +/// +/// The arguments are the three thread coordinates followed by the first two +/// block extents. Pure and host-testable; [`block_linear_thread_id_here`] +/// reads the special registers. +#[must_use] +#[inline(always)] +pub(crate) const fn block_linear_thread_id( + thread_idx_x: u32, + thread_idx_y: u32, + thread_idx_z: u32, + block_dim_x: u32, + block_dim_y: u32, +) -> u32 { + thread_idx_x + block_dim_x * (thread_idx_y + block_dim_y * thread_idx_z) +} + +/// Block-linear id of the calling thread; see [`block_linear_thread_id`]. +/// +/// Reads `threadIdx` and `blockDim`, which exist only on the device. Host +/// tests call [`block_linear_thread_id`] directly. +#[inline(always)] +pub(crate) fn block_linear_thread_id_here() -> u32 { + block_linear_thread_id( + threadIdx_x(), + threadIdx_y(), + threadIdx_z(), + blockDim_x(), + blockDim_y(), + ) +} + // ============================================================================= // Synchronization Intrinsics // ============================================================================= @@ -1395,8 +1441,9 @@ pub fn __unchecked_indexing_config() { /// ``` /// /// The pass currently recognizes explicit counted `while` loops and range `for` -/// loops (`for i in a..b`). Other iterator loops, such as `a..=b` or -/// `.step_by(k)`, are not yet recognized. +/// loops (`for i in a..b` and `(a..b).rev()`) over any integer counter, +/// including `isize`. Other iterator loops, such as `a..=b` or `.step_by(k)`, +/// are not yet recognized. /// /// Loops with several `continue` paths are supported. Full `#[unroll]` also /// preserves `break` paths and multiple exit targets. Partial `#[unroll(N)]` @@ -1429,3 +1476,82 @@ const fn validate_unroll_factor(factor: u32) { "partial unroll factor must be in 2..=1024, or 0 for full unrolling" ); } + +#[cfg(test)] +mod tests { + use super::block_linear_thread_id; + + /// Number of threads in a `(dx, dy, dz)` block whose block-linear id is + /// `selected`. `ManagedBarrier::init_by` and `inval_by` make this + /// comparison on the device, one thread at a time. + fn threads_selected_by(block_dim: (u32, u32, u32), selected: u32) -> u32 { + let (dx, dy, dz) = block_dim; + let mut hits = 0; + for tz in 0..dz { + for ty in 0..dy { + for tx in 0..dx { + if block_linear_thread_id(tx, ty, tz, dx, dy) == selected { + hits += 1; + } + } + } + } + hits + } + + #[test] + fn a_1d_block_numbers_threads_by_x() { + for tx in 0..256 { + assert_eq!(block_linear_thread_id(tx, 0, 0, 256, 1), tx); + } + } + + #[test] + fn a_2d_block_counts_x_fastest_then_y() { + // (32, 4, 1): row y starts at linear id 32 * y, so each row is one warp. + for ty in 0..4 { + for tx in 0..32 { + assert_eq!(block_linear_thread_id(tx, ty, 0, 32, 4), tx + 32 * ty); + } + } + assert_eq!(block_linear_thread_id(0, 1, 0, 32, 4), 32); + assert_eq!(block_linear_thread_id(31, 3, 0, 32, 4), 127); + } + + #[test] + fn a_3d_block_visits_every_id_once_in_x_y_z_order() { + // (8, 4, 2): 64 threads. Nested loops over z, then y, then x must + // produce 0, 1, 2, ... with no gap and no repeat. + let (dx, dy, dz) = (8, 4, 2); + let mut expected = 0; + for tz in 0..dz { + for ty in 0..dy { + for tx in 0..dx { + assert_eq!(block_linear_thread_id(tx, ty, tz, dx, dy), expected); + expected += 1; + } + } + } + assert_eq!(expected, dx * dy * dz); + // One step along z skips a whole x-y plane. + assert_eq!(block_linear_thread_id(0, 0, 1, dx, dy), dx * dy); + } + + #[test] + fn every_in_range_id_selects_one_thread() { + for block_dim in [(64, 1, 1), (32, 4, 1), (8, 4, 2), (16, 2, 2), (24, 3, 1)] { + let (dx, dy, dz) = block_dim; + let threads = dx * dy * dz; + for selected in 0..threads { + assert_eq!( + threads_selected_by(block_dim, selected), + 1, + "block {block_dim:?} id {selected}" + ); + } + // An id at or past the thread count selects nobody, which is why + // init_by and inval_by require the selector to be in range. + assert_eq!(threads_selected_by(block_dim, threads), 0); + } + } +} diff --git a/cuda-oxide/crates/cuda-device/src/tma.rs b/cuda-oxide/crates/cuda-device/src/tma.rs index 3532cab7e4..bf5d72bdef 100644 --- a/cuda-oxide/crates/cuda-device/src/tma.rs +++ b/cuda-oxide/crates/cuda-device/src/tma.rs @@ -113,6 +113,7 @@ //! - **sm_120 (Blackwell)**: Latest TMA capabilities use crate::barrier::Barrier; +use crate::shared::SharedPtr; // ============================================================================= // TMA Descriptor Type @@ -126,8 +127,10 @@ use crate::barrier::Barrier; /// /// # 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 /// @@ -142,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 50a7db4a63..edb513e6c3 100644 --- a/cuda-oxide/crates/cuda-device/src/vector.rs +++ b/cuda-oxide/crates/cuda-device/src/vector.rs @@ -165,31 +165,52 @@ macro_rules! define_vector { define_vector!( /// Two `f32` in one 64-bit transaction. - F32x2, f32, 2, 8 + F32x2, + f32, + 2, + 8 ); define_vector!( /// Four `f32` in one 128-bit transaction. The `float4` equivalent. - F32x4, f32, 4, 16 + F32x4, + f32, + 4, + 16 ); define_vector!( /// Two `f64` in one 128-bit transaction. - F64x2, f64, 2, 16 + F64x2, + f64, + 2, + 16 ); define_vector!( /// Two `u32` in one 64-bit transaction. - U32x2, u32, 2, 8 + U32x2, + u32, + 2, + 8 ); define_vector!( /// Four `u32` in one 128-bit transaction. - U32x4, u32, 4, 16 + U32x4, + u32, + 4, + 16 ); define_vector!( /// Four `u16` in one 64-bit transaction. - U16x4, u16, 4, 8 + U16x4, + u16, + 4, + 8 ); define_vector!( /// Eight `u16` in one 128-bit transaction. The width packed f16 pairs use. - U16x8, u16, 8, 16 + U16x8, + u16, + 8, + 16 ); /// View a flat slice as over-aligned vectors, or `None` if it does not divide. diff --git a/cuda-oxide/crates/cuda-device/src/warp.rs b/cuda-oxide/crates/cuda-device/src/warp.rs index 138b0fce4d..f4043e128c 100644 --- a/cuda-oxide/crates/cuda-device/src/warp.rs +++ b/cuda-oxide/crates/cuda-device/src/warp.rs @@ -209,13 +209,60 @@ pub fn active_mask() -> u32 { /// Computes: `threadIdx.x / 32` /// /// This is a derived value, not a hardware register. -/// Only valid for 1D thread blocks; for multi-dimensional blocks, -/// compute your own warp ID from the linearized thread index. +/// Only valid for 1D thread blocks. In a 2D or 3D block hardware warps are +/// formed from consecutive block-linear thread ids, so use +/// `cooperative_groups::this_thread_block().thread_rank() / 32` there. That +/// is the block-linear warp id [`TmemGuard`](crate::tcgen05::TmemGuard) uses. #[inline(always)] pub fn warp_id() -> u32 { crate::thread::threadIdx_x() / 32 } +/// Warp index of one thread inside its block, in linear thread order. +/// +/// This is `linear_tid / 32`, with `linear_tid` from +/// [`crate::thread::block_linear_thread_id`]: `x` fastest, then `y`, then `z`. +/// +/// For a 1D block this equals `thread_idx_x / 32`, which is what [`warp_id`] +/// returns. For a 2D or 3D block it names the single group of 32 consecutive +/// linear thread ids, not every thread that shares the same `threadIdx.x`. +/// +/// [`warp_id`] stays x-only. Callers that must pick one warp of the block +/// (tensor-memory alloc and dealloc, the cooperative-groups block reductions) +/// use this function instead. +#[must_use] +#[inline(always)] +pub(crate) const fn block_linear_warp_id( + thread_idx_x: u32, + thread_idx_y: u32, + thread_idx_z: u32, + block_dim_x: u32, + block_dim_y: u32, +) -> u32 { + crate::thread::block_linear_thread_id( + thread_idx_x, + thread_idx_y, + thread_idx_z, + block_dim_x, + block_dim_y, + ) / 32 +} + +/// Block-linear warp id of the calling thread. +/// +/// Reads `threadIdx` and `blockDim`. Host tests exercise [`block_linear_warp_id`] +/// directly because those special registers are not available on the host. +#[inline(always)] +pub(crate) fn block_linear_warp_id_here() -> u32 { + block_linear_warp_id( + crate::thread::threadIdx_x(), + crate::thread::threadIdx_y(), + crate::thread::threadIdx_z(), + crate::thread::blockDim_x(), + crate::thread::blockDim_y(), + ) +} + // ============================================================================= // Masked sync intrinsics — operand convention // ============================================================================= @@ -1657,4 +1704,126 @@ mod tests { ); } } + + fn count_matching( + block_dim_x: u32, + block_dim_y: u32, + block_dim_z: u32, + warp: u32, + pred: impl Fn(u32, u32, u32) -> bool, + ) -> u32 { + let mut count = 0u32; + for tz in 0..block_dim_z { + for ty in 0..block_dim_y { + for tx in 0..block_dim_x { + let id = block_linear_warp_id(tx, ty, tz, block_dim_x, block_dim_y); + if id == warp { + assert!(pred(tx, ty, tz), "thread ({tx}, {ty}, {tz}) in warp {warp}"); + count += 1; + } + } + } + } + count + } + + #[test] + fn block_linear_warp_id_is_thread_idx_x_div_32_for_a_1d_block() { + // 256 threads, one dimension: warp n is exactly threads 32n..32n+31. + let block_dim_x = 256; + for n in 0..8u32 { + let count = count_matching(block_dim_x, 1, 1, n, |tx, ty, tz| { + ty == 0 && tz == 0 && tx / 32 == n && (32 * n..32 * n + 32).contains(&tx) + }); + assert_eq!(count, 32, "1D warp {n}"); + } + } + + #[test] + fn block_linear_warp_id_2d_32x4_names_one_row_per_warp() { + // 32x4: linear order is x-fastest, so each y row is its own warp. + // `threadIdx.x / 32` would mark every row as warp 0. + let block_dim_x = 32; + let block_dim_y = 4; + let warp0 = count_matching(block_dim_x, block_dim_y, 1, 0, |tx, ty, tz| { + (0..32).contains(&tx) && ty == 0 && tz == 0 + }); + assert_eq!(warp0, 32, "warp 0 must be only the y=0 row"); + + let warp1 = count_matching(block_dim_x, block_dim_y, 1, 1, |tx, ty, tz| { + (0..32).contains(&tx) && ty == 1 && tz == 0 + }); + assert_eq!(warp1, 32, "warp 1 must be only the y=1 row"); + + for warp in 0..4u32 { + let count = count_matching(block_dim_x, block_dim_y, 1, warp, |tx, ty, tz| { + (0..32).contains(&tx) && ty == warp && tz == 0 + }); + assert_eq!(count, 32, "warp {warp} must contain exactly 32 threads"); + } + } + + #[test] + fn block_linear_warp_id_3d_smoke_partitions_the_block() { + // 8x4x2 = 64 threads → two warps of 32, x fastest. + let (dx, dy, dz) = (8, 4, 2); + let warp0 = count_matching(dx, dy, dz, 0, |tx, ty, tz| { + let linear = tx + dx * (ty + dy * tz); + (0..32).contains(&linear) + }); + let warp1 = count_matching(dx, dy, dz, 1, |tx, ty, tz| { + let linear = tx + dx * (ty + dy * tz); + (32..64).contains(&linear) + }); + assert_eq!(warp0, 32); + assert_eq!(warp1, 32); + assert_eq!(warp0 + warp1, dx * dy * dz); + } + + /// Threads of a block that the pre-fix rule `threadIdx.x / 32 == warp` + /// selects: their count, and how many of them fall into each block-linear + /// warp (index = block-linear warp id; 32 entries cover 1024 threads). + fn x_only_selection(dx: u32, dy: u32, dz: u32, warp: u32) -> (u32, [u32; 32]) { + let mut count = 0u32; + let mut per_warp = [0u32; 32]; + for tz in 0..dz { + for ty in 0..dy { + for tx in 0..dx { + if tx / 32 == warp { + count += 1; + per_warp[block_linear_warp_id(tx, ty, tz, dx, dy) as usize] += 1; + } + } + } + } + (count, per_warp) + } + + #[test] + fn the_x_only_rule_selects_all_four_warps_of_a_32x4_block() { + // Control: every thread of a (32, 4, 1) block has threadIdx.x / 32 == 0, + // so the old selection covers 128 threads in four full warps. The + // block-linear rule selects the 32 threads of warp 0. + let (old_count, per_warp) = x_only_selection(32, 4, 1, 0); + assert_eq!(old_count, 128); + assert_eq!(per_warp[..4], [32, 32, 32, 32]); + assert_eq!(count_matching(32, 4, 1, 0, |_, _, _| true), 32); + } + + #[test] + fn the_x_only_rule_straddles_two_warps_when_block_dim_x_is_48() { + // Control: in a (48, 2, 1) block, target warp 1, the old rule selects + // threadIdx.x in 32..48 of both rows: linear ids 32..=47 and 80..=95, + // half of block-linear warp 1 and half of warp 2. No full warp is + // selected, so a .sync.aligned instruction there is undefined. + let (old_count, per_warp) = x_only_selection(48, 2, 1, 1); + assert_eq!(old_count, 32); + assert_eq!(per_warp[..3], [0, 16, 16]); + // The block-linear rule selects linear ids 32..=63: the end of row 0 + // and the start of row 1, one full hardware warp. + let new_count = count_matching(48, 2, 1, 1, |tx, ty, _| { + (ty == 0 && tx >= 32) || (ty == 1 && tx < 16) + }); + assert_eq!(new_count, 32); + } } 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 f6bf308fe7..7d921a8bb6 100644 --- a/cuda-oxide/crates/cuda-host/README.md +++ b/cuda-oxide/crates/cuda-host/README.md @@ -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/NVlabs/cutile-rs/tree/main/cuda-core) -- shared CUDA driver API crate; `DeviceBuffer`, `simt::LaunchConfig` -- [cuda-async](https://github.com/NVlabs/cutile-rs/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`, `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` diff --git a/cuda-oxide/crates/cuda-host/src/lib.rs b/cuda-oxide/crates/cuda-host/src/lib.rs index 5e6a866a7a..2e5505cdbe 100644 --- a/cuda-oxide/crates/cuda-host/src/lib.rs +++ b/cuda-oxide/crates/cuda-host/src/lib.rs @@ -10,10 +10,10 @@ // 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)] diff --git a/cuda-oxide/crates/cuda-host/src/ltoir.rs b/cuda-oxide/crates/cuda-host/src/ltoir.rs index 810c56a396..5eaad7ef43 100644 --- a/cuda-oxide/crates/cuda-host/src/ltoir.rs +++ b/cuda-oxide/crates/cuda-host/src/ltoir.rs @@ -33,8 +33,8 @@ //! - **libNVVM**: `LIBNVVM_PATH`, then `/nvvm/lib64/libnvvm.so` for //! `` in `CUDA_TOOLKIT_PATH`, `CUDA_HOME`, `CUDA_PATH`, //! `/usr/local/cuda`, `/opt/cuda`, then the system loader. -//! - **nvJitLink**: the same order, using `LIBNVJITLINK_PATH` and -//! `/lib64/libnvJitLink.so`. +//! - **nvJitLink**: the same order, using `LIBNVJITLINK_PATH`, then +//! `/lib64/libnvJitLink.so` and `/lib/libnvJitLink.so`. //! - **libdevice**: `CUDA_OXIDE_LIBDEVICE` env var, then //! `/nvvm/libdevice/libdevice.10.bc` for the same roots. //! - **Artifact target and options**: the emitted `.target` records the diff --git a/cuda-oxide/crates/cuda-host/src/type_id.rs b/cuda-oxide/crates/cuda-host/src/type_id.rs index 8a46e85cf1..6d480a4bf4 100644 --- a/cuda-oxide/crates/cuda-host/src/type_id.rs +++ b/cuda-oxide/crates/cuda-host/src/type_id.rs @@ -23,7 +23,7 @@ //! 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. 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/contracts/scalar.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/scalar.rs index 257a1a85e5..f3da03ecca 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/scalar.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/scalar.rs @@ -311,7 +311,7 @@ pub enum ExtendedMinMaxAdapter { /// Closed contract for converting between scalar and packed values. /// /// The source may be a scalar pair (`f32x2`, two operands) or an already-packed -/// 16-bit or 32-bit register (`f16x2`, `e4m3x2`, `e5m2x2`, one operand), so the +/// 16-bit or 32-bit register (`bf16x2`, `f16x2`, `e4m3x2`, `e5m2x2`, one operand), so the /// operand arity follows [`PackedConversionSourceFormat`] rather than being /// fixed at two. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] @@ -327,6 +327,7 @@ pub struct PackedConversion { #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum PackedConversionSourceFormat { + Bf16x2, E4m3x2, E5m2x2, F16x2, @@ -341,13 +342,14 @@ impl PackedConversionSourceFormat { pub fn operand_count(self) -> usize { match self { Self::F32x2 => 2, - Self::E4m3x2 | Self::E5m2x2 | Self::F16x2 => 1, + Self::Bf16x2 | Self::E4m3x2 | Self::E5m2x2 | Self::F16x2 => 1, } } /// PTX source-type token, used as the trailing `cvt` modifier. pub fn ptx_token(self) -> &'static str { match self { + Self::Bf16x2 => "bf16x2", Self::E4m3x2 => "e4m3x2", Self::E5m2x2 => "e5m2x2", Self::F16x2 => "f16x2", diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/tma.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/tma.rs index d4690333ec..5270e609c7 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/tma.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/contracts/tma.rs @@ -30,12 +30,17 @@ impl Tma { #[serde(rename_all = "snake_case")] pub enum TmaOperation { G2sTile1d, + G2sCtaTile1d, G2sTile2d, + G2sCtaTile2d, G2sTile2dMulticast, G2sTile2dMulticastCg2, G2sTile3d, + G2sCtaTile3d, G2sTile4d, + G2sCtaTile4d, G2sTile5d, + G2sCtaTile5d, S2gTile1d, S2gTile2d, S2gTile3d, @@ -82,16 +87,28 @@ pub enum TmaOperation { } impl TmaOperation { + pub const fn is_cta_g2s(self) -> bool { + matches!( + self, + Self::G2sCtaTile1d + | Self::G2sCtaTile2d + | Self::G2sCtaTile3d + | Self::G2sCtaTile4d + | Self::G2sCtaTile5d + ) + } + pub const fn dimensions(self) -> Option { match self { - Self::G2sTile1d | Self::S2gTile1d => Some(1), + Self::G2sTile1d | Self::G2sCtaTile1d | Self::S2gTile1d => Some(1), Self::G2sTile2d + | Self::G2sCtaTile2d | Self::G2sTile2dMulticast | Self::G2sTile2dMulticastCg2 | Self::S2gTile2d => Some(2), - Self::G2sTile3d | Self::S2gTile3d => Some(3), - Self::G2sTile4d | Self::S2gTile4d => Some(4), - Self::G2sTile5d | Self::S2gTile5d => Some(5), + Self::G2sTile3d | Self::G2sCtaTile3d | Self::S2gTile3d => Some(3), + Self::G2sTile4d | Self::G2sCtaTile4d | Self::S2gTile4d => Some(4), + Self::G2sTile5d | Self::G2sCtaTile5d | Self::S2gTile5d => Some(5), Self::Reduce | Self::CommitGroup | Self::WaitGroup 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 081f4b33b7..be1b534f97 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/model/overlay.rs @@ -81,6 +81,8 @@ pub struct OverlayShardFile { #[serde(default)] pub packed_conversion_fp8: Option, #[serde(default)] + pub packed_conversion_fp8_bf16x2: Option, + #[serde(default)] pub packed_conversion_fp8_f16x2: Option, #[serde(default)] pub scalar_conversion: Option, @@ -166,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, @@ -322,6 +326,12 @@ pub struct ClcAdmissionVariant { pub struct TmaAdmission { pub llvm_evidence_profile: String, pub libnvvm_evidence_profile: String, + /// Evidence profiles for the CTA-local G2S variants. A shard that lists + /// those variants must name both; an older shard without them parses. + #[serde(default)] + pub cta_llvm_evidence_profile: Option, + #[serde(default)] + pub cta_libnvvm_evidence_profile: Option, #[serde(default)] pub reduce_llvm_evidence_profile: Option, #[serde(default)] @@ -341,6 +351,16 @@ pub struct TmaAdmissionVariant { pub operation: TmaOperation, } +impl TmaAdmission { + /// Whether the shard lists the CTA-local G2S variants. Those variants + /// carry their own evidence profiles and need overlay shard schema 65. + pub fn admits_cta_g2s(&self) -> bool { + self.variants + .iter() + .any(|variant| variant.operation.is_cta_g2s()) + } +} + /// One reviewed tensor-reduction operation and its reserved ABI ID. #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] @@ -505,6 +525,22 @@ pub enum PackedConversionFp8Format { E5m2x2, } +/// Compact admission for packing `bf16x2` into `e4m3x2` or `e5m2x2`. +/// +/// Pack only. The pinned LLVM import has no unscaled unpack record for +/// `bf16x2`; its reverse records are the block-scaled `scale_n2_ue8m0` forms. +/// The PTX ISA 9.2 unpack forms are a later admission. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PackedConversionFp8Bf16x2Admission { + pub llvm_evidence_profile: String, + pub libnvvm_evidence_profile: String, + pub runtime_validation: RuntimeValidation, + pub fp8_formats: Vec, + pub relu_variants: bool, + pub product_count: usize, +} + /// Which way an FP8/`f16x2` conversion runs. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize)] #[serde(rename_all = "snake_case")] 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 ecc48feca7..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, @@ -913,12 +915,17 @@ pub(super) fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String { let is_g2s = matches!( operation, TmaOperation::G2sTile1d + | TmaOperation::G2sCtaTile1d | TmaOperation::G2sTile2d + | TmaOperation::G2sCtaTile2d | TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 | TmaOperation::G2sTile3d + | TmaOperation::G2sCtaTile3d | TmaOperation::G2sTile4d + | TmaOperation::G2sCtaTile4d | TmaOperation::G2sTile5d + | TmaOperation::G2sCtaTile5d ); let is_s2g = matches!( operation, @@ -935,6 +942,11 @@ pub(super) fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String { output.push_str( "/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes.\n", ); + if operation.is_cta_g2s() { + output.push_str( + "/// `barrier` must point into the issuing CTA's shared memory, as for the other\n/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type.\n", + ); + } } else if is_s2g { output.push_str( "/// `src` and `tensor_map` must remain valid until the committed copy group completes.\n", @@ -961,8 +973,15 @@ pub(super) fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String { output.push_str("#[inline(never)]\n"); if is_g2s { let dimensions = dimensions.unwrap(); + // The CTA-local forms take a typed CTA-shared destination; the + // raw ABI keeps `*mut u8`, as `*mut Barrier` keeps `*mut u64`. + let destination = if operation.is_cta_g2s() { + "dst: SharedPtr" + } else { + "dst: *mut u8" + }; let mut arguments = vec![ - "dst: *mut u8".to_owned(), + destination.to_owned(), "tensor_map: *const TmaDescriptor".to_owned(), ]; let mut values = vec!["dst".to_owned(), "tensor_map".to_owned()]; @@ -1857,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 @@ -1869,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) {{", @@ -1888,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 {{", @@ -1899,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(); @@ -2092,7 +2124,8 @@ pub(super) fn render_compat_packed_conversion( // follows the record's own argument types rather than a fixed f32 pair. let parameter_names: Vec<&str> = match packed_conversion_source(record) { PackedConversionSourceFormat::F32x2 => vec![argument_names.0, argument_names.1], - PackedConversionSourceFormat::E4m3x2 + PackedConversionSourceFormat::Bf16x2 + | PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 | PackedConversionSourceFormat::F16x2 => vec!["packed"], }; diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/cluster_tensor.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/cluster_tensor.rs index 0a66d7d299..12e0cbcccb 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/cluster_tensor.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/cluster_tensor.rs @@ -28,7 +28,6 @@ use pliron::{ builtin::op_interfaces::{NOpdsInterface, NResultsInterface}, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -124,7 +123,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -300,7 +298,6 @@ use pliron::{ builtin::{op_interfaces::{NOpdsInterface, NResultsInterface}, types::IntegerType}, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -407,7 +404,6 @@ pub(in crate::render) fn render_dialect_debug_control(catalog: &CatalogFile, has common_traits::Verify,\n\ context::{Context, Ptr},\n\ identifier::Identifier,\n\ - location::Located,\n\ op::Op,\n\ operation::Operation,\n\ result::Error,\n\ @@ -501,7 +497,7 @@ pub(in crate::render) fn render_dialect_clc(catalog: &CatalogFile, hash: &str) - assert_eq!(clc_intrinsics(catalog).count(), 6); let mut output = rust_header(catalog, hash); output.push_str( - "//! Generated Cluster Launch Control operations.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::{TypeHandle, Typed},\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", + "//! Generated Cluster Launch Control operations.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::{TypeHandle, Typed},\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", ); for record in clc_intrinsics(catalog) { let result_count = record.dialect.results.len(); @@ -547,7 +543,7 @@ pub(in crate::render) fn render_dialect_execution_control( assert_eq!(execution_controls(catalog).count(), 8); let mut output = rust_header(catalog, hash); output.push_str( - "//! Generated counted-barrier, grid-dependency, and register-control operations.\n\nuse pliron::{\n builtin::{attributes::IntegerAttr, op_interfaces::{NOpdsInterface, NResultsInterface}, types::{IntegerType, Signedness}},\n common_traits::Verify,\n context::{Context, Ptr},\n identifier::Identifier,\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n verify_err,\n};\nuse pliron::utils::apint::APInt;\nuse pliron_derive::pliron_op;\nuse std::num::NonZeroUsize;\n\n", + "//! Generated counted-barrier, grid-dependency, and register-control operations.\n\nuse pliron::{\n builtin::{attributes::IntegerAttr, op_interfaces::{NOpdsInterface, NResultsInterface}, types::{IntegerType, Signedness}},\n common_traits::Verify,\n context::{Context, Ptr},\n identifier::Identifier,\n op::Op,\n operation::Operation,\n result::Error,\n verify_err,\n};\nuse pliron::utils::apint::APInt;\nuse pliron_derive::pliron_op;\nuse std::num::NonZeroUsize;\n\n", ); for record in execution_controls(catalog).filter(|record| record.family != "register_control") { let operand_count = record.dialect.operands.len(); @@ -630,7 +626,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 9e76a78491..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 @@ -22,7 +22,7 @@ pub(in crate::render) fn render_dialect_movmatrix(catalog: &CatalogFile, hash: & } let record = movmatrix(catalog).next().unwrap(); output.push_str( - "//! Structural operation for generated in-register matrix transpose.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + "//! Structural operation for generated in-register matrix transpose.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", ); writeln!(output, "/// {}", record.summary).unwrap(); writeln!( @@ -66,7 +66,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -396,7 +395,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -499,7 +497,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -518,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)] @@ -776,7 +773,6 @@ use pliron::{ builtin::{attributes::IntegerAttr, ops::ConstantOp, types::{FP32Type, IntegerType, Signedness}}, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/scalar_packed.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/scalar_packed.rs index 045f7c5c83..036c53c8b8 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/scalar_packed.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/dialect/scalar_packed.rs @@ -26,7 +26,7 @@ pub(in crate::render) fn render_dialect_integer_minmax( ) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operations for generated extended integer min/max.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + "//! Structural operations for generated extended integer min/max.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", ); for record in integer_minmaxes(catalog) { let signedness = if record.rust.result == "i32" { @@ -91,7 +91,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -282,7 +281,7 @@ pub(super) fn register(ctx: &mut Context) { pub(in crate::render) fn render_dialect_packed_alu(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operations for generated packed floating-point arithmetic.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(\n ctx: &Context,\n ty: pliron::r#type::TypeHandle,\n width: u32,\n) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", + "//! Structural operations for generated packed floating-point arithmetic.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(\n ctx: &Context,\n ty: pliron::r#type::TypeHandle,\n width: u32,\n) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", ); for record in packed_alus(catalog) { let arity = record.rust.arguments.len(); @@ -344,7 +343,7 @@ pub(in crate::render) fn render_dialect_packed_conversion( ) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operation for generated packed conversion.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\nfn is_integer_width(\n ctx: &Context,\n ty: pliron::r#type::TypeHandle,\n width: u32,\n) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", + "//! Structural operation for generated packed conversion.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\nfn is_integer_width(\n ctx: &Context,\n ty: pliron::r#type::TypeHandle,\n width: u32,\n) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", ); for record in packed_conversions(catalog) { let result_width = packed_conversion_result_width(record); @@ -383,7 +382,8 @@ pub(in crate::render) fn render_dialect_packed_conversion( ) .unwrap(); } - PackedConversionSourceFormat::E4m3x2 + PackedConversionSourceFormat::Bf16x2 + | PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 | PackedConversionSourceFormat::F16x2 => { let source_width = packed_conversion_source_width(record); @@ -443,7 +443,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -584,7 +583,6 @@ use pliron::{ builtin::{op_interfaces::NResultsInterface, types::{FP32Type, FP64Type}}, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -754,7 +752,6 @@ use pliron::{ builtin::{op_interfaces::{NOpdsInterface, NResultsInterface}, types::{FP32Type, FP64Type, IntegerType, Signedness}}, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -913,7 +910,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -1088,7 +1084,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 217452d806..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 @@ -19,7 +19,7 @@ use std::fmt::Write as _; pub(in crate::render) fn render_dialect_elect(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Generated warp leader-election operation.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::{TypeHandle, Typed},\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", + "//! Generated warp leader-election operation.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::{TypeHandle, Typed},\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", ); for record in elect_intrinsics(catalog) { writeln!(output, "/// {}", record.summary).unwrap(); @@ -50,7 +50,7 @@ pub(in crate::render) fn render_dialect_elect(catalog: &CatalogFile, hash: &str) pub(in crate::render) fn render_dialect_sreg(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural NVVM operations for generated special-register reads.\n\nuse pliron::{\n builtin::op_interfaces::{NOpdsInterface, NResultsInterface},\n builtin::types::IntegerType,\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", + "//! Structural NVVM operations for generated special-register reads.\n\nuse pliron::{\n builtin::op_interfaces::{NOpdsInterface, NResultsInterface},\n builtin::types::IntegerType,\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", ); for record in sregs(catalog) { let width = record.scalar_width().unwrap(); @@ -123,7 +123,7 @@ pub(in crate::render) fn render_dialect_sync(catalog: &CatalogFile, hash: &str) pub(in crate::render) fn render_dialect_vote(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operations for the generated `vote.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", + "//! Structural operations for the generated `vote.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", ); for record in vote_intrinsics(catalog) { let result_width = match record.vote.as_ref().unwrap().mode { @@ -178,7 +178,7 @@ pub(in crate::render) fn render_dialect_vote(catalog: &CatalogFile, hash: &str) pub(in crate::render) fn render_dialect_active_mask(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operation for the generated active warp mask.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", + "//! Structural operation for the generated active warp mask.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\n", ); for record in active_masks(catalog) { debug_assert_eq!( @@ -227,7 +227,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -478,7 +477,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -676,7 +674,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -787,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", @@ -800,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", @@ -808,6 +807,7 @@ fn verify_mbarrier_extended( ), MbarrierExtendedOperation::FenceProxyAsyncSharedCta | MbarrierExtendedOperation::FenceMbarrierInitReleaseCluster + | MbarrierExtendedOperation::FenceAcquireSharedClusterCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericReleaseSharedCtaCluster | MbarrierExtendedOperation::FenceProxyAsyncGenericAcquireSharedClusterCluster => ( 0, @@ -853,7 +853,7 @@ fn verify_mbarrier_extended( pub(in crate::render) fn render_dialect_warp_match(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operations for the generated `match.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", + "//! Structural operations for the generated `match.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_integer_width(ctx: &Context, ty: pliron::r#type::TypeHandle, width: u32) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == width)\n}\n\n", ); for record in warp_matches(catalog) { let warp_match = record.warp_match.as_ref().unwrap(); @@ -904,7 +904,7 @@ pub(in crate::render) fn render_dialect_warp_match(catalog: &CatalogFile, hash: pub(in crate::render) fn render_dialect_warp_barrier(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operation for generated warp synchronization.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + "//! Structural operation for generated warp synchronization.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::IntegerType,\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", ); for record in warp_barriers(catalog) { debug_assert_eq!( @@ -947,7 +947,7 @@ pub(in crate::render) fn render_dialect_warp_barrier(catalog: &CatalogFile, hash pub(in crate::render) fn render_dialect_warp_shuffle(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operations for the generated `shfl.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\nfn is_i64(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 64)\n}\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\n", + "//! Structural operations for the generated `shfl.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\nfn is_i64(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 64)\n}\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\n", ); for record in warp_shuffles(catalog) { let shuffle = record.warp_shuffle.as_ref().unwrap(); @@ -1024,7 +1024,7 @@ pub(in crate::render) fn render_dialect_warp_shuffle(catalog: &CatalogFile, hash pub(in crate::render) fn render_dialect_dotprod(catalog: &CatalogFile, hash: &str) -> String { let mut output = rust_header(catalog, hash); output.push_str( - "//! Structural operations for generated packed integer dot products.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + "//! Structural operations for generated packed integer dot products.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", ); for record in dot_products(catalog) { let signedness = match record.rust.result.as_str() { @@ -1080,11 +1080,11 @@ pub(in crate::render) fn render_dialect_redux(catalog: &CatalogFile, hash: &str) let has_f32 = redux(catalog).any(|record| record.dialect.results == ["f32"]); if has_f32 { output.push_str( - "//! Structural NVVM operations for the closed generated `redux.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\n", + "//! Structural NVVM operations for the closed generated `redux.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{FP32Type, IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\nfn is_f32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx).downcast_ref::().is_some()\n}\n\n", ); } else { output.push_str( - "//! Structural NVVM operations for the closed generated `redux.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n location::Located,\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", + "//! Structural NVVM operations for the closed generated `redux.sync` family.\n\nuse pliron::{\n builtin::{\n op_interfaces::{NOpdsInterface, NResultsInterface},\n types::{IntegerType, Signedness},\n },\n common_traits::Verify,\n context::{Context, Ptr},\n op::Op,\n operation::Operation,\n result::Error,\n r#type::Typed,\n value::Value,\n verify_err,\n};\nuse pliron_derive::pliron_op;\n\nfn is_i32(ctx: &Context, ty: pliron::r#type::TypeHandle) -> bool {\n ty.deref(ctx)\n .downcast_ref::()\n .is_some_and(|integer| integer.width() == 32)\n}\n\n", ); } for record in redux(catalog) { 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 f2de275c66..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", @@ -366,6 +367,12 @@ pub(in crate::render) const SPARSE_MMA_ORDERED_TF32_METADATA_RULE: &str = pub(in crate::render) fn sparse_mma_metadata_rule(mma: &SparseMma) -> &'static str { match (mma.metadata, mma.a_element) { + // PTX 8.4 requires these six ordered index pairs for plain FP8 mma.sp + // too; this does not change the instruction to mma.sp::ordered_metadata. + // https://docs.nvidia.com/cuda/archive/12.4.1/parallel-thread-execution/index.html#warp-level-sparse-matrix-storage + (SparseMmaMetadata::Standard, SparseMmaElement::E4m3 | SparseMmaElement::E5m2) => { + SPARSE_MMA_ORDERED_METADATA_RULE + } (SparseMmaMetadata::Standard, _) => SPARSE_MMA_STANDARD_METADATA_RULE, (SparseMmaMetadata::Ordered, SparseMmaElement::Tf32) => { SPARSE_MMA_ORDERED_TF32_METADATA_RULE diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/families/packed.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/families/packed.rs index f4977f10af..f377c9a597 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/families/packed.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/families/packed.rs @@ -240,7 +240,9 @@ pub(in crate::render) fn packed_conversion_is_closed_recipe( let adapter_matches = conversion.adapter == match conversion.source_format { Src::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, - Src::E4m3x2 | Src::E5m2x2 | Src::F16x2 => PackedConversionAdapter::Identity, + Src::Bf16x2 | Src::E4m3x2 | Src::E5m2x2 | Src::F16x2 => { + PackedConversionAdapter::Identity + } }; adapter_matches @@ -285,6 +287,20 @@ pub(in crate::render) fn packed_conversion_is_closed_recipe( Round::NearestEven, Sat::SatfiniteRelu ) + | (Src::Bf16x2, Dst::E4m3x2, Round::NearestEven, Sat::Satfinite) + | ( + Src::Bf16x2, + Dst::E4m3x2, + Round::NearestEven, + Sat::SatfiniteRelu + ) + | (Src::Bf16x2, Dst::E5m2x2, Round::NearestEven, Sat::Satfinite) + | ( + Src::Bf16x2, + Dst::E5m2x2, + Round::NearestEven, + Sat::SatfiniteRelu + ) | (Src::E4m3x2, Dst::F16x2, Round::NearestEven, Sat::None) | (Src::E4m3x2, Dst::F16x2, Round::NearestEven, Sat::Relu) | (Src::E5m2x2, Dst::F16x2, Round::NearestEven, Sat::None) @@ -305,7 +321,7 @@ pub(in crate::render) fn packed_conversion_source( /// Register width of a single packed source operand. pub(in crate::render) fn packed_conversion_source_width(record: &CatalogIntrinsic) -> u32 { match packed_conversion_source(record) { - PackedConversionSourceFormat::F16x2 => 32, + PackedConversionSourceFormat::Bf16x2 | PackedConversionSourceFormat::F16x2 => 32, PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => 16, PackedConversionSourceFormat::F32x2 => { unreachable!("f32x2 conversions do not have a single packed source") @@ -319,7 +335,9 @@ pub(in crate::render) fn packed_conversion_rust_arguments( ) -> Vec<&'static str> { match packed_conversion_source(record) { PackedConversionSourceFormat::F32x2 => vec!["f32", "f32"], - PackedConversionSourceFormat::F16x2 => vec!["u32"], + PackedConversionSourceFormat::Bf16x2 | PackedConversionSourceFormat::F16x2 => { + vec!["u32"] + } PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => vec!["u16"], } } @@ -330,7 +348,9 @@ pub(in crate::render) fn packed_conversion_dialect_operands( ) -> Vec<&'static str> { match packed_conversion_source(record) { PackedConversionSourceFormat::F32x2 => vec!["f32", "f32"], - PackedConversionSourceFormat::F16x2 => vec!["i32"], + PackedConversionSourceFormat::Bf16x2 | PackedConversionSourceFormat::F16x2 => { + vec!["i32"] + } PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 => vec!["i16"], } } @@ -372,7 +392,7 @@ pub(in crate::render) fn packed_conversion_constraint(record: &CatalogIntrinsic) ) { (16, PackedConversionSourceFormat::F32x2) => "=h,f,f", (32, PackedConversionSourceFormat::F32x2) => "=r,f,f", - (16, PackedConversionSourceFormat::F16x2) => "=h,r", + (16, PackedConversionSourceFormat::Bf16x2 | PackedConversionSourceFormat::F16x2) => "=h,r", (32, PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2) => "=r,h", _ => unreachable!("closed packed-conversion result width and source format"), } @@ -426,7 +446,8 @@ pub(in crate::render) fn packed_conversion_ptx_mnemonic(record: &CatalogIntrinsi conversion.adapter, match conversion.source_format { PackedConversionSourceFormat::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, - PackedConversionSourceFormat::E4m3x2 + PackedConversionSourceFormat::Bf16x2 + | PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 | PackedConversionSourceFormat::F16x2 => PackedConversionAdapter::Identity, } diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/importer.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/importer.rs index 78478f8c59..26883fe979 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/importer.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/importer.rs @@ -2402,6 +2402,18 @@ fn tma_arms(catalog: &CatalogFile) -> String { ) .unwrap(); } + TmaOperation::G2sCtaTile1d + | TmaOperation::G2sCtaTile2d + | TmaOperation::G2sCtaTile3d + | TmaOperation::G2sCtaTile4d + | TmaOperation::G2sCtaTile5d => { + writeln!( + output, + " Ok(Some(super::super::tma::emit_tma_g2s_cta(\n ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, {}, {marker:?},\n )?))", + operation.dimensions().unwrap() + ) + .unwrap(); + } TmaOperation::G2sTile2dMulticast => { writeln!( output, @@ -3183,9 +3195,6 @@ fn importer_shard_imports( output.push_str("use pliron::input_err;\n"); } let mut location_items = Vec::new(); - if body.contains(".set_loc(") { - location_items.push("Located"); - } if dispatch_signature || uses_identifier(body, "Location") { location_items.push("Location"); } diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/lowering.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/lowering.rs index 36b73243f1..09c0b9f030 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/lowering.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/lowering.rs @@ -1085,7 +1085,8 @@ fn packed_conversion_impls(catalog: &CatalogFile) -> String { ) .unwrap(); } - PackedConversionSourceFormat::E4m3x2 + PackedConversionSourceFormat::Bf16x2 + | PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 | PackedConversionSourceFormat::F16x2 => { debug_assert_eq!(conversion.adapter, PackedConversionAdapter::Identity); @@ -1519,6 +1520,18 @@ fn tma_impls(catalog: &CatalogFile) -> String { ) .unwrap(); } + TmaOperation::G2sCtaTile1d + | TmaOperation::G2sCtaTile2d + | TmaOperation::G2sCtaTile3d + | TmaOperation::G2sCtaTile4d + | TmaOperation::G2sCtaTile5d => { + writeln!( + output, + " convert_g2s_cta(ctx, rewriter, self.get_operation(), operands_info, {})", + operation.dimensions().unwrap() + ) + .unwrap(); + } TmaOperation::G2sTile2dMulticast => output.push_str( " convert_g2s(ctx, rewriter, self.get_operation(), operands_info, 2, true)\n", ), @@ -2099,6 +2112,7 @@ const LOWERING_INTRINSIC_HELPERS: &[(&str, &str)] = &[ ("ReduceConfig", "tma::ReduceConfig"), ("convert_control", "tma::convert_control"), ("convert_g2s", "tma::convert_g2s"), + ("convert_g2s_cta", "tma::convert_g2s_cta"), ( "convert_g2s_multicast_cg2", "tma::convert_g2s_multicast_cg2", @@ -2237,9 +2251,6 @@ fn lowering_shard_imports(catalog: &CatalogFile, body: &str, shard_file: bool) - { output.push_str("use pliron::irbuild::rewriter::Rewriter;\n"); } - if body.contains(".loc()") || body.contains(".set_loc(") { - output.push_str("use pliron::location::Located;\n"); - } if body.contains("::get_concrete_op_info(") || body.contains(".get_operation()") { output.push_str("use pliron::op::Op;\n"); } diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/probes.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/probes.rs index 3c7858866c..980c034896 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/probes.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/probes.rs @@ -84,12 +84,17 @@ fn render_tma_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str let is_g2s = matches!( operation, TmaOperation::G2sTile1d + | TmaOperation::G2sCtaTile1d | TmaOperation::G2sTile2d + | TmaOperation::G2sCtaTile2d | TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 | TmaOperation::G2sTile3d + | TmaOperation::G2sCtaTile3d | TmaOperation::G2sTile4d + | TmaOperation::G2sCtaTile4d | TmaOperation::G2sTile5d + | TmaOperation::G2sCtaTile5d ); let coordinates = std::iter::repeat_n("i32", dimensions) .collect::>() @@ -103,6 +108,14 @@ fn render_tma_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str .collect::>() .join(", "); if is_g2s { + let cta = matches!( + operation, + TmaOperation::G2sCtaTile1d + | TmaOperation::G2sCtaTile2d + | TmaOperation::G2sCtaTile3d + | TmaOperation::G2sCtaTile4d + | TmaOperation::G2sCtaTile5d + ); let multicast = matches!( operation, TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 @@ -112,11 +125,19 @@ fn render_tma_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str } else { format!(", {coordinates}") }; - writeln!( - output, - "declare void @{symbol}(ptr addrspace(7), ptr addrspace(3), ptr{declaration_coordinates}, i16, i64, i1, i1, i32) #0\n" - ) - .unwrap(); + if cta { + writeln!( + output, + "declare void @{symbol}(ptr addrspace(3), ptr addrspace(3), ptr{declaration_coordinates}, i64, i1) #0\n" + ) + .unwrap(); + } else { + writeln!( + output, + "declare void @{symbol}(ptr addrspace(7), ptr addrspace(3), ptr{declaration_coordinates}, i16, i64, i1, i1, i32) #0\n" + ) + .unwrap(); + } let parameters = if coordinate_parameters.is_empty() { String::new() } else { @@ -129,9 +150,15 @@ fn render_tma_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str record.id ) .unwrap(); - output.push_str( - " %dst = addrspacecast ptr %dst_generic to ptr addrspace(7)\n %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", - ); + if cta { + output.push_str( + " %dst = addrspacecast ptr %dst_generic to ptr addrspace(3)\n %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + } else { + output.push_str( + " %dst = addrspacecast ptr %dst_generic to ptr addrspace(7)\n %barrier = addrspacecast ptr %barrier_generic to ptr addrspace(3)\n", + ); + } let arguments = if coordinate_arguments.is_empty() { String::new() } else { @@ -143,11 +170,19 @@ fn render_tma_probe(catalog: &CatalogFile, record: &CatalogIntrinsic, hash: &str } else { 0 }; - writeln!( - output, - " call void @{symbol}(ptr addrspace(7) %dst, ptr addrspace(3) %barrier, ptr %tensor_map{arguments}, i16 {mask}, i64 0, i1 {multicast}, i1 false, i32 {group}) #0" - ) - .unwrap(); + if cta { + writeln!( + output, + " call void @{symbol}(ptr addrspace(3) %dst, ptr addrspace(3) %barrier, ptr %tensor_map{arguments}, i64 0, i1 false) #0" + ) + .unwrap(); + } else { + writeln!( + output, + " call void @{symbol}(ptr addrspace(7) %dst, ptr addrspace(3) %barrier, ptr %tensor_map{arguments}, i16 {mask}, i64 0, i1 {multicast}, i1 false, i32 {group}) #0" + ) + .unwrap(); + } } else { let declaration_coordinates = if coordinates.is_empty() { String::new() 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 6477e238cf..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 Result + { + output.push_str( + "/// The destination and barrier must lie in the issuing CTA's own shared memory, and the tensor map must be a live descriptor for this dimensionality.\n\ + /// A generic destination or barrier pointer is converted with `cvta.to.shared`, so a cluster-shared pointer from `map_shared_rank` would yield a wrong local offset; copies into another CTA's shared memory use the cluster form.\n\ + /// Keep every object alive until the asynchronous copy completes. Only the designated issuing thread may start this transfer.\n", + ) + } TmaAdapter::G2sPointersCoordinatesBarrierInjectDefaults | TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults => output .push_str( diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/reference.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/reference.rs index 781d43b5ee..133f459a5c 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/reference.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/reference.rs @@ -7,9 +7,9 @@ use crate::model::{ BackendLoweringMechanism, CatalogFile, ClusterBarrierOrdering, ClusterMemorySourceContract, CpAsyncMbarrierOperation, CpAsyncMbarrierStateSpace, EvidenceArtifactKind, IntrinsicBackend, MbarrierBasicOperation, MbarrierExtendedSourceContract, PackedAtomicFormat, - PackedConversionRounding, PackedConversionSaturation, RegisterMmaAdapter, RegisterMmaOperation, - RegisterMmaOverflow, RuntimeValidation, SparseMmaOverflow, WarpMatchAdapter, - WarpShuffleValueKind, WgmmaControlMode, + PackedConversionRounding, PackedConversionSaturation, RegisterMmaAdapter, RegisterMmaKind, + RegisterMmaOperation, RegisterMmaOverflow, RuntimeValidation, SparseMmaOverflow, + WarpMatchAdapter, WarpShuffleValueKind, WgmmaControlMode, }; use crate::render::common::{ backend_label, evidence_stage_label, hardware_target_label, llvm, lowering_mechanism_label, @@ -120,9 +120,14 @@ pub(super) fn render_reference(catalog: &CatalogFile, hash: &str) -> 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 67123ff4e6..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(), 1029); + 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")); @@ -791,6 +823,38 @@ fn register_mma_kind_is_explicit_and_legacy_ir_stays_compatible() { assert!(targets.contains("kind_matches && accumulator_matches")); } +#[test] +fn sparse_fp8_metadata_contract_requires_ordered_indices_without_changing_integer_rules() { + let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let catalog = crate::resolve::resolve(&repo_root).unwrap(); + for id in [ + "mma_sp_m16n8k64_f32_e4m3_e4m3_f32", + "mma_sp_m16n8k64_f32_e4m3_e5m2_f32", + "mma_sp_m16n8k64_f32_e5m2_e4m3_f32", + "mma_sp_m16n8k64_f32_e5m2_e5m2_f32", + ] { + let record = sparse_mmas(&catalog) + .find(|record| record.id == id) + .unwrap(); + assert_eq!( + sparse_mma_metadata_rule(record.sparse_mma.as_ref().unwrap()), + "Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior.", + "{id}" + ); + assert!(sparse_mma_template(record).starts_with("mma.sp.sync.aligned.")); + } + for id in ["mma_sp_m16n8k64_s32_s8", "mma_sp_m16n8k64_s32_s4"] { + let record = sparse_mmas(&catalog) + .find(|record| record.id == id) + .unwrap(); + assert_eq!( + sparse_mma_metadata_rule(record.sparse_mma.as_ref().unwrap()), + "Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior.", + "{id}" + ); + } +} + #[test] fn sparse_mma_rendering_enforces_selector_and_keeps_family_distinct() { let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); @@ -1072,8 +1136,8 @@ fn sparse_mma_rendering_enforces_selector_and_keeps_family_distinct() { assert!(targets.contains("GeneratedSparseMmaOverflow::NotApplicable")); assert!(targets.contains("GeneratedHardwareAlternative::ExactArchitecture(120)")); - assert_eq!(raw.matches(SPARSE_MMA_STANDARD_METADATA_RULE).count(), 36); - assert_eq!(raw.matches(SPARSE_MMA_ORDERED_METADATA_RULE).count(), 88); + assert_eq!(raw.matches(SPARSE_MMA_STANDARD_METADATA_RULE).count(), 32); + assert_eq!(raw.matches(SPARSE_MMA_ORDERED_METADATA_RULE).count(), 92); assert_eq!( raw.matches(SPARSE_MMA_ORDERED_TF32_METADATA_RULE).count(), 2 @@ -1082,13 +1146,13 @@ fn sparse_mma_rendering_enforces_selector_and_keeps_family_distinct() { compatibility .matches(SPARSE_MMA_STANDARD_METADATA_RULE) .count(), - 36 + 32 ); assert_eq!( compatibility .matches(SPARSE_MMA_ORDERED_METADATA_RULE) .count(), - 88 + 92 ); assert_eq!( compatibility @@ -1146,11 +1210,11 @@ fn sparse_mma_rendering_enforces_selector_and_keeps_family_distinct() { assert!(reference.contains("LLVM source record uses A, B, C, metadata, selector order")); assert_eq!( reference.matches(SPARSE_MMA_STANDARD_METADATA_RULE).count(), - 36 + 32 ); assert_eq!( reference.matches(SPARSE_MMA_ORDERED_METADATA_RULE).count(), - 88 + 92 ); assert_eq!( reference diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mod.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mod.rs index 729db95fa8..9afc2769f3 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mod.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/mod.rs @@ -42,7 +42,7 @@ fn catalog_with_clc() -> CatalogFile { fn catalog_with_tma() -> CatalogFile { let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); let catalog = crate::resolve::test_catalog_with_tma(&repo_root).unwrap(); - assert_eq!(tma_intrinsics(&catalog).count(), 111); + assert_eq!(tma_intrinsics(&catalog).count(), 116); catalog } diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/scalar_misc.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/scalar_misc.rs index b947253880..6c555c0456 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/scalar_misc.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/scalar_misc.rs @@ -108,7 +108,7 @@ fn packed_alu_and_conversion_render_exact_pure_inline_ptx_adapters() { let catalog = crate::resolve::resolve(&repo_root).unwrap(); validate_renderable(&catalog).unwrap(); assert_eq!(packed_alus(&catalog).count(), 30); - assert_eq!(packed_conversions(&catalog).count(), 18); + assert_eq!(packed_conversions(&catalog).count(), 22); let dialect = render_dialect_packed_alu(&catalog, "test-hash"); for op in [ diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs index d431282c51..dec130798d 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs @@ -419,7 +419,13 @@ fn tma_rendering_preserves_api_and_injects_backend_defaults() { assert!(compatibility.contains( "pub unsafe fn cp_async_bulk_tensor_reduce_xor_im2col_5d(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, coord2: i32, coord3: i32, coord4: i32)" )); + assert!(compatibility.contains( + "pub unsafe fn cp_async_bulk_tensor_2d_g2s_cta(dst: SharedPtr, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, barrier: *mut Barrier)" + )); for dimensions in 1..=5 { + assert!(compatibility.contains(&format!( + "pub unsafe fn cp_async_bulk_tensor_{dimensions}d_g2s_cta(dst: SharedPtr," + ))); assert!(compatibility.contains(&format!( "pub unsafe fn cp_async_bulk_prefetch_tensor_{dimensions}d_l2(" ))); @@ -437,8 +443,8 @@ fn tma_rendering_preserves_api_and_injects_backend_defaults() { assert!(compatibility.contains("pub fn fence_proxy_tensormap_generic_release_system()")); let dialect = render_dialect_tma(&catalog, "test-hash"); - assert_eq!(dialect.matches("pub struct ").count(), 111); - assert_eq!(dialect.matches("NResultsInterface<0>").count(), 111); + assert_eq!(dialect.matches("pub struct ").count(), 116); + assert_eq!(dialect.matches("NResultsInterface<0>").count(), 116); assert!(dialect.contains("NOpdsInterface<10>")); assert!(dialect.contains("CpAsyncBulkWaitGroupReadOp::register(ctx)")); assert!(dialect.contains("CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2Op::register(ctx)")); 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 f78c55a7b2..62a5fd3e9a 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/render/validate.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/render/validate.rs @@ -733,10 +733,15 @@ pub(super) fn validate_renderable(catalog: &CatalogFile) -> Result<()> { && match (tma.operation, tma.adapter) { ( TmaOperation::G2sTile1d + | TmaOperation::G2sCtaTile1d | TmaOperation::G2sTile2d + | TmaOperation::G2sCtaTile2d | TmaOperation::G2sTile3d + | TmaOperation::G2sCtaTile3d | TmaOperation::G2sTile4d - | TmaOperation::G2sTile5d, + | TmaOperation::G2sCtaTile4d + | TmaOperation::G2sTile5d + | TmaOperation::G2sCtaTile5d, TmaAdapter::G2sPointersCoordinatesBarrierInjectDefaults, ) => { !record.rust.safe @@ -1105,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, @@ -1135,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/driver.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/driver.rs index 41a224c375..782dfd196f 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/driver.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/driver.rs @@ -243,6 +243,8 @@ pub(crate) fn test_catalog_with_tma(repo_root: &Path) -> Result { let admission = TmaAdmission { llvm_evidence_profile: "llvm-tma-test".into(), libnvvm_evidence_profile: "libnvvm-tma-test".into(), + cta_llvm_evidence_profile: Some("llvm-tma-test".into()), + cta_libnvvm_evidence_profile: Some("libnvvm-tma-test".into()), reduce_llvm_evidence_profile: Some("llvm-tma-reduce-test".into()), reduce_libnvvm_evidence_profile: Some("libnvvm-tma-reduce-test".into()), runtime_validation: RuntimeValidation::Unexecuted, 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/register.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/register.rs index 55b69af7ee..0dd3ad3d17 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/register.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/mma/register.rs @@ -800,8 +800,13 @@ pub(in crate::resolve) fn register_mma_recipe(mma: &RegisterMma) -> Option Option { 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/families/packed_conversion.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/packed_conversion.rs index b3c0a82a7d..ee75422523 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/packed_conversion.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/packed_conversion.rs @@ -6,7 +6,8 @@ use crate::model::{ BackendLoweringMechanism, ImportedIntrinsic, IntrinsicBackend, IntrinsicSource, OverlayBackendLowering, OverlayIntrinsic, PackedConversionAdapter, - PackedConversionDestinationFormat, PackedConversionFp8Admission, PackedConversionFp8Direction, + PackedConversionDestinationFormat, PackedConversionFp8Admission, + PackedConversionFp8Bf16x2Admission, PackedConversionFp8Direction, PackedConversionFp8F16x2Admission, PackedConversionFp8Format, PackedConversionRounding, PackedConversionSaturation, PackedConversionSourceFormat, RuntimeValidation, }; @@ -16,6 +17,17 @@ use std::collections::BTreeSet; use crate::resolve::guards::*; +// PTX ISA 9.1 lists `cvt.rn.satfinite{.relu}.{e4m3x2,e5m2x2}.bf16x2` for +// sm_100f, sm_110f and sm_120f "or higher in the same family". The ISA's +// `.target` section lets family-specific features run on f-targets and on +// a-targets of the same family, nvcc defines `__CUDA_ARCH_FAMILY_SPECIFIC__` +// for `compute_100a` as well, and ptxas 13.4 and the pinned llc accept every +// member below while rejecting the unsuffixed targets. cuda-oxide-codegen +// matches a family target only on the literal `f` suffix, so the a-targets +// are enumerated here, as the other Blackwell family records do. +pub(in crate::resolve) const PACKED_CONVERSION_BF16X2_TARGETS: &str = + "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f"; + pub(in crate::resolve) struct PackedConversionRecipe { pub(in crate::resolve) id: &'static str, pub(in crate::resolve) abi_id: &'static str, @@ -35,12 +47,107 @@ pub(in crate::resolve) fn packed_conversion_recipe( ) -> Option { match conversion.source_format { PackedConversionSourceFormat::F32x2 => packed_conversion_recipe_f32x2(conversion), + PackedConversionSourceFormat::Bf16x2 => packed_conversion_recipe_fp8_bf16x2(conversion), PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 | PackedConversionSourceFormat::F16x2 => packed_conversion_recipe_fp8_f16x2(conversion), } } +/// Recipes for packing `bf16x2` into `e4m3x2` or `e5m2x2`. +/// +/// Like the `f16x2` packs, these lower through inline PTX: the MIR side carries +/// both halves in one integer register while the typed intrinsics are declared +/// over `<2 x bfloat>` (see [`packed_conversion_uses_typed_nvvm`]). The +/// recorded evidence assembles the typed and inline routes to the same cubin. +/// PTX ISA 9.1 introduces the instruction on the Blackwell families named by +/// [`PACKED_CONVERSION_BF16X2_TARGETS`], not on the Ada `sm_89` floor used by +/// `f16x2`. +pub(in crate::resolve) fn packed_conversion_recipe_fp8_bf16x2( + conversion: &crate::model::PackedConversion, +) -> Option { + match ( + conversion.source_format, + conversion.destination_format, + conversion.rounding, + conversion.saturation, + ) { + ( + PackedConversionSourceFormat::Bf16x2, + PackedConversionDestinationFormat::E4m3x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Satfinite, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_e4m3x2_bf16x2", + abi_id: "i1030", + operation_key: "packed.convert.bf16x2.e4m3x2.nearest_even.satfinite", + rust_name: "cvt_rn_satfinite_e4m3x2_bf16x2", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2", + dialect_op_type: "CvtRnSatfiniteE4m3x2Bf16x2Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_e4m3x2_bf16x2", + source_record: "int_nvvm_bf16x2_to_e4m3x2_rn_satfinite", + llvm_symbol: "llvm.nvvm.bf16x2.to.e4m3x2.rn.satfinite", + llvm_result: "i16", + summary: "Converts packed bf16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order.", + }), + ( + PackedConversionSourceFormat::Bf16x2, + PackedConversionDestinationFormat::E4m3x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::SatfiniteRelu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_relu_e4m3x2_bf16x2", + abi_id: "i1031", + operation_key: "packed.convert.bf16x2.e4m3x2.nearest_even.satfinite.relu", + rust_name: "cvt_rn_satfinite_relu_e4m3x2_bf16x2", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2", + dialect_op_type: "CvtRnSatfiniteReluE4m3x2Bf16x2Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e4m3x2_bf16x2", + source_record: "int_nvvm_bf16x2_to_e4m3x2_rn_relu_satfinite", + llvm_symbol: "llvm.nvvm.bf16x2.to.e4m3x2.rn.relu.satfinite", + llvm_result: "i16", + summary: "Converts packed bf16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order.", + }), + ( + PackedConversionSourceFormat::Bf16x2, + PackedConversionDestinationFormat::E5m2x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::Satfinite, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_e5m2x2_bf16x2", + abi_id: "i1032", + operation_key: "packed.convert.bf16x2.e5m2x2.nearest_even.satfinite", + rust_name: "cvt_rn_satfinite_e5m2x2_bf16x2", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2", + dialect_op_type: "CvtRnSatfiniteE5m2x2Bf16x2Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_e5m2x2_bf16x2", + source_record: "int_nvvm_bf16x2_to_e5m2x2_rn_satfinite", + llvm_symbol: "llvm.nvvm.bf16x2.to.e5m2x2.rn.satfinite", + llvm_result: "i16", + summary: "Converts packed bf16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order.", + }), + ( + PackedConversionSourceFormat::Bf16x2, + PackedConversionDestinationFormat::E5m2x2, + PackedConversionRounding::NearestEven, + PackedConversionSaturation::SatfiniteRelu, + ) => Some(PackedConversionRecipe { + id: "cvt_rn_satfinite_relu_e5m2x2_bf16x2", + abi_id: "i1033", + operation_key: "packed.convert.bf16x2.e5m2x2.nearest_even.satfinite.relu", + rust_name: "cvt_rn_satfinite_relu_e5m2x2_bf16x2", + compatibility_path: "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2", + dialect_op_type: "CvtRnSatfiniteReluE5m2x2Bf16x2Op", + dialect_op_name: "nvvm.cvt_rn_satfinite_relu_e5m2x2_bf16x2", + source_record: "int_nvvm_bf16x2_to_e5m2x2_rn_relu_satfinite", + llvm_symbol: "llvm.nvvm.bf16x2.to.e5m2x2.rn.relu.satfinite", + llvm_result: "i16", + summary: "Converts packed bf16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order.", + }), + _ => None, + } +} + /// Recipes for the packed FP8 conversions whose other side is `f16x2`. /// /// Keyed on source as well as destination: unpacking to `f16x2` shares its @@ -434,6 +541,11 @@ pub(in crate::resolve) fn packed_conversion_source_types( vec!["f32".into(), "f32".into()], vec!["f32".into(), "f32".into()], ), + PackedConversionSourceFormat::Bf16x2 => ( + vec!["u32".into()], + vec!["i32".into()], + vec!["v2bf16".into()], + ), PackedConversionSourceFormat::F16x2 => { (vec!["u32".into()], vec!["i32".into()], vec!["v2f16".into()]) } @@ -454,18 +566,23 @@ pub(in crate::resolve) fn packed_conversion_result_width( pub(in crate::resolve) fn packed_conversion_floor( conversion: &crate::model::PackedConversion, -) -> (&'static str, &'static str) { +) -> (&'static str, Option<&'static str>, &'static str) { + // `bf16x2` to packed FP8 is a PTX ISA 9.1 family-specific conversion. + // It is not available on the Ada `sm_89` floor used by `f16x2`. + if conversion.source_format == PackedConversionSourceFormat::Bf16x2 { + return ("9.1", None, PACKED_CONVERSION_BF16X2_TARGETS); + } // FP8 on either side carries the Ada floor, including when FP8 is the // source and the destination is the older `f16x2`. if packed_conversion_uses_fp8(conversion) { - return ("8.1", "sm_89"); + return ("8.1", Some("sm_89"), "all"); } match conversion.destination_format { PackedConversionDestinationFormat::Bf16x2 | PackedConversionDestinationFormat::F16x2 => { - ("7.0", "sm_80") + ("7.0", Some("sm_80"), "all") } PackedConversionDestinationFormat::E4m3x2 | PackedConversionDestinationFormat::E5m2x2 => { - ("8.1", "sm_89") + ("8.1", Some("sm_89"), "all") } } } @@ -482,12 +599,13 @@ pub(in crate::resolve) fn packed_conversion_backend_mechanism( /// Whether the conversion lowers through a typed NVVM intrinsic call. /// -/// Only the scalar-`f32` pair does. The `f16x2` sources and destinations carry -/// their halves in one integer register on the MIR side, while the typed -/// intrinsics are declared over `<2 x half>`; routing them through inline PTX -/// avoids a packed-half MIR type for no change in emitted code. The recorded -/// evidence assembles both routes to a byte-identical cubin, so the typed route -/// can be switched on later without changing the PTX contract. +/// Only the scalar-`f32` pair does. The `f16x2` and `bf16x2` sources and the +/// `f16x2` destinations carry their halves in one integer register on the MIR +/// side, while the typed intrinsics are declared over `<2 x half>` or +/// `<2 x bfloat>`; routing them through inline PTX avoids a packed-half MIR +/// type for no change in emitted code. The recorded evidence assembles both +/// routes to a byte-identical cubin, so the typed route can be switched on +/// later without changing the PTX contract. pub(in crate::resolve) fn packed_conversion_uses_typed_nvvm( conversion: &crate::model::PackedConversion, ) -> bool { @@ -563,6 +681,64 @@ pub(in crate::resolve) fn expand_packed_conversion_fp8_admission( Ok(records) } +pub(in crate::resolve) fn expand_packed_conversion_fp8_bf16x2_admission( + admission: &PackedConversionFp8Bf16x2Admission, +) -> Result> { + ensure!( + admission.runtime_validation == RuntimeValidation::Unexecuted, + "FP8 bf16x2 conversion runtime validation may be marked executed only with GPU evidence" + ); + ensure!( + admission.fp8_formats + == [ + PackedConversionFp8Format::E4m3x2, + PackedConversionFp8Format::E5m2x2, + ], + "compact FP8 bf16x2 conversion admission must list the canonical two formats" + ); + ensure!( + admission.relu_variants, + "compact FP8 bf16x2 conversion admission must admit the ReLU variants" + ); + ensure!( + admission.product_count + == admission + .fp8_formats + .len() + .checked_mul(2) + .context("compact FP8 bf16x2 conversion product count overflow")? + && admission.product_count == 4, + "compact FP8 bf16x2 conversion product_count must be exactly 4" + ); + + let mut records = Vec::with_capacity(admission.product_count); + for &fp8_format in &admission.fp8_formats { + for relu in [false, true] { + let conversion = crate::model::PackedConversion { + source_format: PackedConversionSourceFormat::Bf16x2, + destination_format: match fp8_format { + PackedConversionFp8Format::E4m3x2 => PackedConversionDestinationFormat::E4m3x2, + PackedConversionFp8Format::E5m2x2 => PackedConversionDestinationFormat::E5m2x2, + }, + rounding: PackedConversionRounding::NearestEven, + saturation: if relu { + PackedConversionSaturation::SatfiniteRelu + } else { + PackedConversionSaturation::Satfinite + }, + adapter: PackedConversionAdapter::Identity, + }; + records.push(packed_conversion_overlay_record( + conversion, + &admission.llvm_evidence_profile, + &admission.libnvvm_evidence_profile, + )?); + } + } + ensure!(records.len() == admission.product_count); + Ok(records) +} + pub(in crate::resolve) fn expand_packed_conversion_fp8_f16x2_admission( admission: &PackedConversionFp8F16x2Admission, ) -> Result> { @@ -669,7 +845,7 @@ pub(in crate::resolve) fn packed_conversion_overlay_record( let result_width = packed_conversion_result_width(&conversion); let rust_result = format!("u{result_width}"); let dialect_result = format!("i{result_width}"); - let (minimum_ptx, minimum_sm) = packed_conversion_floor(&conversion); + let (minimum_ptx, minimum_sm, targets) = packed_conversion_floor(&conversion); let (rust_arguments, dialect_operands, llvm_arguments) = packed_conversion_source_types(&conversion); // `cvt` writes one destination and reads every source operand. @@ -707,9 +883,9 @@ pub(in crate::resolve) fn packed_conversion_overlay_record( convergent: false, execution_scope: "thread".into(), minimum_ptx: minimum_ptx.into(), - minimum_sm: Some(minimum_sm.into()), + minimum_sm: minimum_sm.map(str::to_owned), ptx_result: rust_result, - targets: "all".into(), + targets: targets.into(), ptx_isa_version: "9.3".into(), ptx_isa_section: "9.7.9.22 Data Movement and Conversion Instructions: cvt".into(), ptx_isa_url: "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt".into(), @@ -725,7 +901,7 @@ pub(in crate::resolve) fn packed_conversion_overlay_record( evidence_profile: evidence_profile.into(), targets: None, minimum_ptx: Some(minimum_ptx.into()), - minimum_sm: Some(minimum_sm.into()), + minimum_sm: minimum_sm.map(str::to_owned), }) .collect(), packed_atomic: None, @@ -790,7 +966,8 @@ pub(in crate::resolve) fn validate_packed_conversion_policy( // packed operand is forwarded unchanged. let expected_adapter = match conversion.source_format { PackedConversionSourceFormat::F32x2 => PackedConversionAdapter::ReverseHighLowOperands, - PackedConversionSourceFormat::E4m3x2 + PackedConversionSourceFormat::Bf16x2 + | PackedConversionSourceFormat::E4m3x2 | PackedConversionSourceFormat::E5m2x2 | PackedConversionSourceFormat::F16x2 => PackedConversionAdapter::Identity, }; @@ -808,7 +985,7 @@ pub(in crate::resolve) fn validate_packed_conversion_policy( let result_width = packed_conversion_result_width(conversion); let rust_result = format!("u{result_width}"); let dialect_result = format!("i{result_width}"); - let (minimum_ptx, minimum_sm) = packed_conversion_floor(conversion); + let (minimum_ptx, minimum_sm, targets) = packed_conversion_floor(conversion); let (expected_rust_arguments, expected_dialect_operands, expected_llvm_arguments) = packed_conversion_source_types(conversion); ensure!( @@ -864,9 +1041,9 @@ pub(in crate::resolve) fn validate_packed_conversion_policy( && !policy.convergent && policy.execution_scope == "thread" && policy.minimum_ptx == minimum_ptx - && policy.minimum_sm.as_deref() == Some(minimum_sm) + && policy.minimum_sm.as_deref() == minimum_sm && policy.ptx_result == rust_result - && policy.targets == "all" + && policy.targets == targets && policy.ptx_isa_version == "9.3" && policy.ptx_isa_section == "9.7.9.22 Data Movement and Conversion Instructions: cvt" && policy.ptx_isa_url @@ -916,7 +1093,7 @@ pub(in crate::resolve) fn validate_packed_conversion_policy( for lowering in &policy.backend_lowerings { ensure!( lowering.minimum_ptx.as_deref() == Some(minimum_ptx) - && lowering.minimum_sm.as_deref() == Some(minimum_sm) + && lowering.minimum_sm.as_deref() == minimum_sm && !lowering.evidence_profile.trim().is_empty(), "{} backend {:?} does not carry its exact packed-conversion floor", policy.id, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/tma.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/tma.rs index 99c0ca90a0..c0172af859 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/tma.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/families/tma.rs @@ -19,14 +19,19 @@ use crate::resolve::guards::*; pub(in crate::resolve) const TMA_BLACKWELL_TARGETS: &str = "sm_100a|sm_101a|sm_103a|sm_110a"; pub(in crate::resolve) const TENSOR_MAP_REPLACE_TARGETS: &str = "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f|sm_90a"; -pub(in crate::resolve) const TMA_OPERATIONS: [TmaOperation; 47] = [ +pub(in crate::resolve) const TMA_OPERATIONS: [TmaOperation; 52] = [ TmaOperation::G2sTile1d, + TmaOperation::G2sCtaTile1d, TmaOperation::G2sTile2d, + TmaOperation::G2sCtaTile2d, TmaOperation::G2sTile2dMulticast, TmaOperation::G2sTile2dMulticastCg2, TmaOperation::G2sTile3d, + TmaOperation::G2sCtaTile3d, TmaOperation::G2sTile4d, + TmaOperation::G2sCtaTile4d, TmaOperation::G2sTile5d, + TmaOperation::G2sCtaTile5d, TmaOperation::S2gTile1d, TmaOperation::S2gTile2d, TmaOperation::S2gTile3d, @@ -409,6 +414,15 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { "CpAsyncBulkTensorG2sTile1dOp", "nvvm.cp_async_bulk_tensor_g2s_tile_1d", ), + TmaOperation::G2sCtaTile1d => ( + "i1034", + "cp_async_bulk_tensor_1d_g2s_cta", + "memory.copy.async.bulk.tensor.g2s.cta.tile.1d", + "int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_1d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.1d", + "CpAsyncBulkTensorG2sCtaTile1dOp", + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_1d", + ), TmaOperation::G2sTile2d => ( "i0329", "cp_async_bulk_tensor_2d_g2s", @@ -418,6 +432,15 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { "CpAsyncBulkTensorG2sTile2dOp", "nvvm.cp_async_bulk_tensor_g2s_tile_2d", ), + TmaOperation::G2sCtaTile2d => ( + "i1035", + "cp_async_bulk_tensor_2d_g2s_cta", + "memory.copy.async.bulk.tensor.g2s.cta.tile.2d", + "int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_2d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.2d", + "CpAsyncBulkTensorG2sCtaTile2dOp", + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_2d", + ), TmaOperation::G2sTile2dMulticast => ( "i0330", "cp_async_bulk_tensor_2d_g2s_multicast", @@ -445,6 +468,15 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { "CpAsyncBulkTensorG2sTile3dOp", "nvvm.cp_async_bulk_tensor_g2s_tile_3d", ), + TmaOperation::G2sCtaTile3d => ( + "i1036", + "cp_async_bulk_tensor_3d_g2s_cta", + "memory.copy.async.bulk.tensor.g2s.cta.tile.3d", + "int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_3d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.3d", + "CpAsyncBulkTensorG2sCtaTile3dOp", + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_3d", + ), TmaOperation::G2sTile4d => ( "i0333", "cp_async_bulk_tensor_4d_g2s", @@ -454,6 +486,15 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { "CpAsyncBulkTensorG2sTile4dOp", "nvvm.cp_async_bulk_tensor_g2s_tile_4d", ), + TmaOperation::G2sCtaTile4d => ( + "i1037", + "cp_async_bulk_tensor_4d_g2s_cta", + "memory.copy.async.bulk.tensor.g2s.cta.tile.4d", + "int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_4d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.4d", + "CpAsyncBulkTensorG2sCtaTile4dOp", + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_4d", + ), TmaOperation::G2sTile5d => ( "i0334", "cp_async_bulk_tensor_5d_g2s", @@ -463,6 +504,15 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { "CpAsyncBulkTensorG2sTile5dOp", "nvvm.cp_async_bulk_tensor_g2s_tile_5d", ), + TmaOperation::G2sCtaTile5d => ( + "i1038", + "cp_async_bulk_tensor_5d_g2s_cta", + "memory.copy.async.bulk.tensor.g2s.cta.tile.5d", + "int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_5d", + "llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.5d", + "CpAsyncBulkTensorG2sCtaTile5dOp", + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_5d", + ), TmaOperation::S2gTile1d => ( "i0335", "cp_async_bulk_tensor_1d_s2g", @@ -829,12 +879,17 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { let is_g2s = matches!( operation, TmaOperation::G2sTile1d + | TmaOperation::G2sCtaTile1d | TmaOperation::G2sTile2d + | TmaOperation::G2sCtaTile2d | TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 | TmaOperation::G2sTile3d + | TmaOperation::G2sCtaTile3d | TmaOperation::G2sTile4d + | TmaOperation::G2sCtaTile4d | TmaOperation::G2sTile5d + | TmaOperation::G2sCtaTile5d ); let is_s2g = matches!( operation, @@ -848,6 +903,7 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { operation, TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 ); + let cta_g2s = operation.is_cta_g2s(); let cg2 = operation == TmaOperation::G2sTile2dMulticastCg2; let prefetch_coordinates = operation.prefetch_coordinate_count(); let is_release_fence = matches!( @@ -877,10 +933,19 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { } dialect_operands.extend(["ptr", "ptr", "ptr"]); dialect_operands.extend(std::iter::repeat_n("i32", dimensions.unwrap())); - dialect_operands.extend(["i16", "i64"]); - llvm_arguments.extend(["shared_cluster_ptr", "shared_ptr", "ptr"]); - llvm_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); - llvm_arguments.extend(["i16", "i64", "i1", "i1", "i32"]); + if !cta_g2s { + dialect_operands.push("i16"); + } + dialect_operands.push("i64"); + if cta_g2s { + llvm_arguments.extend(["shared_ptr", "shared_ptr", "ptr"]); + llvm_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); + llvm_arguments.extend(["i64", "i1"]); + } else { + llvm_arguments.extend(["shared_cluster_ptr", "shared_ptr", "ptr"]); + llvm_arguments.extend(std::iter::repeat_n("i32", dimensions.unwrap())); + llvm_arguments.extend(["i16", "i64", "i1", "i1", "i32"]); + } ( if multicast { TmaAdapter::G2sPointersCoordinatesBarrierMaskInjectDefaults @@ -892,6 +957,8 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { true, if multicast { "Starts a multicast TMA tile copy from global to cluster shared memory." + } else if cta_g2s { + "Starts a TMA tile copy from global to the issuing CTA's shared memory." } else { "Starts a TMA tile copy from global to cluster shared memory." }, @@ -1094,7 +1161,11 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { "bulk".into(), "tensor".into(), format!("{}d", dimensions.unwrap()), - "shared::cluster".into(), + if cta_g2s { + "shared::cta".into() + } else { + "shared::cluster".into() + }, "global".into(), "tile".into(), "mbarrier::complete_tx::bytes".into(), @@ -1308,13 +1379,14 @@ pub(in crate::resolve) fn tma_recipe(operation: TmaOperation) -> TmaRecipe { unreachable!("TMA operation category was matched") }; - let minimum_ptx = if blackwell_tma || operation == TmaOperation::ReplaceSwizzleAtomicity { - "8.6" - } else if is_replace || is_fence { - "8.3" - } else { - "8.0" - }; + let minimum_ptx = + if blackwell_tma || cta_g2s || operation == TmaOperation::ReplaceSwizzleAtomicity { + "8.6" + } else if is_replace || is_fence { + "8.3" + } else { + "8.0" + }; let (minimum_sm, targets) = if blackwell_tma { (None, TMA_BLACKWELL_TARGETS) } else if operation == TmaOperation::ReplaceSwizzleAtomicity { @@ -1389,14 +1461,32 @@ pub(in crate::resolve) fn expand_tma_admission( && !admission.libnvvm_evidence_profile.trim().is_empty(), "compact TMA admission requires both backend evidence profiles" ); + let admits_cta_g2s = admission.admits_cta_g2s(); + if admits_cta_g2s { + ensure!( + admission + .cta_llvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()) + && admission + .cta_libnvvm_evidence_profile + .as_deref() + .is_some_and(|profile| !profile.trim().is_empty()), + "compact TMA admission with CTA-local G2S variants requires cta_llvm_evidence_profile and cta_libnvvm_evidence_profile" + ); + } + let expected_operations = TMA_OPERATIONS + .into_iter() + .filter(|operation| admits_cta_g2s || !operation.is_cta_g2s()) + .collect::>(); ensure!( admission .variants .iter() .map(|variant| variant.operation) - .eq(TMA_OPERATIONS), + .eq(expected_operations.iter().copied()), "compact TMA admission must list all {} operations in canonical order", - TMA_OPERATIONS.len() + expected_operations.len() ); let mut records = admission @@ -1410,6 +1500,22 @@ pub(in crate::resolve) fn expand_tma_admission( recipe.id, recipe.abi_id ); + let llvm_evidence_profile = if recipe.operation.is_cta_g2s() { + admission + .cta_llvm_evidence_profile + .as_ref() + .expect("validated TMA CTA-local LLVM evidence profile") + } else { + &admission.llvm_evidence_profile + }; + let libnvvm_evidence_profile = if recipe.operation.is_cta_g2s() { + admission + .cta_libnvvm_evidence_profile + .as_ref() + .expect("validated TMA CTA-local libNVVM evidence profile") + } else { + &admission.libnvvm_evidence_profile + }; Ok(OverlayIntrinsic { id: recipe.id.into(), abi_id: variant.abi_id.clone(), @@ -1462,7 +1568,7 @@ pub(in crate::resolve) fn expand_tma_admission( OverlayBackendLowering { backend: IntrinsicBackend::LlvmNvptx, mechanism: recipe.llvm_mechanism, - evidence_profile: admission.llvm_evidence_profile.clone(), + evidence_profile: llvm_evidence_profile.clone(), targets: None, minimum_ptx: Some(recipe.minimum_ptx.into()), minimum_sm: recipe.minimum_sm.map(Into::into), @@ -1470,7 +1576,7 @@ pub(in crate::resolve) fn expand_tma_admission( OverlayBackendLowering { backend: IntrinsicBackend::LibNvvm, mechanism: BackendLoweringMechanism::InlinePtx, - evidence_profile: admission.libnvvm_evidence_profile.clone(), + evidence_profile: libnvvm_evidence_profile.clone(), targets: None, minimum_ptx: Some(recipe.minimum_ptx.into()), minimum_sm: recipe.minimum_sm.map(Into::into), @@ -1781,6 +1887,14 @@ pub(in crate::resolve) fn validate_tma_reduction_policy( pub(in crate::resolve) fn tma_imported_properties(operation: TmaOperation) -> Vec { assert_ne!(operation, TmaOperation::Reduce); let dimensions = operation.dimensions(); + if operation.is_cta_g2s() { + return vec![ + format!("ImmArg", dimensions.unwrap() + 4), + "IntrConvergent".into(), + "ReadOnly".into(), + "WriteOnly".into(), + ]; + } if matches!( operation, TmaOperation::G2sTile1d diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/guards.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/guards.rs index 4575d19c80..ac42753c76 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/guards.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/guards.rs @@ -307,10 +307,18 @@ pub(super) fn selection_matches_tma_policy( "TMA_G2S_TILE_CG0_1D", "cp.async.bulk.tensor.1d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0}}], [$mbar];", ), + TmaOperation::G2sCtaTile1d => ( + "TMA_G2S_CTA_TILE_1D", + "cp.async.bulk.tensor.1d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0}}], [$mbar];", + ), TmaOperation::G2sTile2d => ( "TMA_G2S_TILE_CG0_2D", "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1}}], [$mbar];", ), + TmaOperation::G2sCtaTile2d => ( + "TMA_G2S_CTA_TILE_2D", + "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1}}], [$mbar];", + ), TmaOperation::G2sTile2dMulticast => ( "TMA_G2S_TILE_CG0_2D_MC", "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster$cg [$dst], [$tmap, {{$d0, $d1}}], [$mbar], $mc;", @@ -323,14 +331,26 @@ pub(super) fn selection_matches_tma_policy( "TMA_G2S_TILE_CG0_3D", "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2}}], [$mbar];", ), + TmaOperation::G2sCtaTile3d => ( + "TMA_G2S_CTA_TILE_3D", + "cp.async.bulk.tensor.3d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2}}], [$mbar];", + ), TmaOperation::G2sTile4d => ( "TMA_G2S_TILE_CG0_4D", "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2, $d3}}], [$mbar];", ), + TmaOperation::G2sCtaTile4d => ( + "TMA_G2S_CTA_TILE_4D", + "cp.async.bulk.tensor.4d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2, $d3}}], [$mbar];", + ), TmaOperation::G2sTile5d => ( "TMA_G2S_TILE_CG0_5D", "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes$cg [$dst], [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], [$mbar];", ), + TmaOperation::G2sCtaTile5d => ( + "TMA_G2S_CTA_TILE_5D", + "cp.async.bulk.tensor.5d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], [$mbar];", + ), TmaOperation::S2gTile1d => ( "TMA_TENSOR_S2G_TILE_1D", "cp.async.bulk.tensor.1d.global.shared::cta.tile.bulk_group [$tmap, {{$d0}}], [$src];", @@ -456,29 +476,48 @@ pub(super) fn selection_matches_tma_policy( if matches!( operation, TmaOperation::G2sTile1d + | TmaOperation::G2sCtaTile1d | TmaOperation::G2sTile2d + | TmaOperation::G2sCtaTile2d | TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 | TmaOperation::G2sTile3d + | TmaOperation::G2sCtaTile3d | TmaOperation::G2sTile4d + | TmaOperation::G2sCtaTile4d | TmaOperation::G2sTile5d + | TmaOperation::G2sCtaTile5d ) { let dimensions = operation.dimensions().unwrap(); - immediate_bindings.push(crate::model::ImportedImmediateBinding { - argument_index: dimensions + 5, - value: if matches!( - operation, - TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 - ) { - -1 - } else { - 0 - }, - }); - immediate_bindings.push(crate::model::ImportedImmediateBinding { - argument_index: dimensions + 6, - value: 0, - }); + if matches!( + operation, + TmaOperation::G2sCtaTile1d + | TmaOperation::G2sCtaTile2d + | TmaOperation::G2sCtaTile3d + | TmaOperation::G2sCtaTile4d + | TmaOperation::G2sCtaTile5d + ) { + immediate_bindings.push(crate::model::ImportedImmediateBinding { + argument_index: dimensions + 4, + value: 0, + }); + } else { + immediate_bindings.push(crate::model::ImportedImmediateBinding { + argument_index: dimensions + 5, + value: if matches!( + operation, + TmaOperation::G2sTile2dMulticast | TmaOperation::G2sTile2dMulticastCg2 + ) { + -1 + } else { + 0 + }, + }); + immediate_bindings.push(crate::model::ImportedImmediateBinding { + argument_index: dimensions + 6, + value: 0, + }); + } } else if matches!( operation, TmaOperation::S2gTile1d diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/overlay.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/overlay.rs index 03a02dbea6..703332219a 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/overlay.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/overlay.rs @@ -19,7 +19,7 @@ use super::guards::*; pub(super) const OVERLAY_SCHEMA: u32 = 44; pub(super) const MINIMUM_OVERLAY_SHARD_SCHEMA: u32 = 26; -pub(super) const OVERLAY_SHARD_SCHEMA: u32 = 64; +pub(super) const OVERLAY_SHARD_SCHEMA: u32 = 65; pub(super) const REGISTER_MMA_F8F6F4_SHARD_SCHEMA: u32 = 46; pub(super) const REGISTER_MMA_F8F6F4_F16_SHARD_SCHEMA: u32 = 47; pub(super) const REGISTER_MMA_MXF8F6F4_SHARD_SCHEMA: u32 = 60; @@ -32,6 +32,7 @@ pub(super) const SPARSE_MMA_FP8_F32_SHARD_SCHEMA: u32 = 64; pub(super) const PRMT_SHARD_SCHEMA: u32 = 28; pub(super) const PACKED_CONVERSION_FP8_SHARD_SCHEMA: u32 = 29; pub(super) const PACKED_CONVERSION_FP8_F16X2_SHARD_SCHEMA: u32 = 59; +pub(super) const PACKED_CONVERSION_FP8_BF16X2_SHARD_SCHEMA: u32 = 65; pub(super) const CLUSTER_SREG_SHARD_SCHEMA: u32 = 30; pub(super) const CLUSTER_BARRIER_SHARD_SCHEMA: u32 = 31; pub(super) const SPECIAL_REGISTER_SHARD_SCHEMA: u32 = 32; @@ -42,6 +43,7 @@ pub(super) const CLUSTER_MEMORY_SHARD_SCHEMA: u32 = 39; pub(super) const CLC_SHARD_SCHEMA: u32 = 40; pub(super) const TMA_SHARD_SCHEMA: u32 = 61; pub(super) const TMA_REDUCTION_SHARD_SCHEMA: u32 = 62; +pub(super) const TMA_CTA_SHARD_SCHEMA: u32 = 65; pub(super) const MBARRIER_EXTENDED_SHARD_SCHEMA: u32 = 40; pub(super) const WGMMA_CONTROL_SHARD_SCHEMA: u32 = 38; pub(super) const TCGEN05_SHARD_SCHEMA: u32 = 42; @@ -114,6 +116,7 @@ pub(super) fn read_overlay( let sparse_mma_ampere_float_admission = shard.sparse_mma_ordered_ampere_float.take(); let prmt_admission = shard.prmt.take(); let packed_conversion_fp8_admission = shard.packed_conversion_fp8.take(); + let packed_conversion_fp8_bf16x2_admission = shard.packed_conversion_fp8_bf16x2.take(); let packed_conversion_fp8_f16x2_admission = shard.packed_conversion_fp8_f16x2.take(); let scalar_conversion_admission = shard.scalar_conversion.take(); let scalar_arithmetic_admission = shard.scalar_arithmetic.take(); @@ -254,6 +257,13 @@ pub(super) fn read_overlay( ); shard.intrinsics = expand_packed_conversion_fp8_admission(&admission)?; } + if let Some(admission) = packed_conversion_fp8_bf16x2_admission { + ensure!( + shard.family == "packed_conversion" && shard.intrinsics.is_empty(), + "compact FP8 bf16x2 conversion admission must be the only content of a packed_conversion shard" + ); + shard.intrinsics = expand_packed_conversion_fp8_bf16x2_admission(&admission)?; + } if let Some(admission) = packed_conversion_fp8_f16x2_admission { ensure!( shard.family == "packed_conversion" && shard.intrinsics.is_empty(), @@ -482,6 +492,12 @@ pub(super) fn validate_overlay_shard_schema_with_max( "compact FP8 f16x2 conversion admission requires overlay shard schema {}", PACKED_CONVERSION_FP8_F16X2_SHARD_SCHEMA ); + ensure!( + shard.packed_conversion_fp8_bf16x2.is_none() + || shard.schema >= PACKED_CONVERSION_FP8_BF16X2_SHARD_SCHEMA, + "compact FP8 bf16x2 conversion admission requires overlay shard schema {}", + PACKED_CONVERSION_FP8_BF16X2_SHARD_SCHEMA + ); ensure!( shard.scalar_conversion.is_none() || shard.schema >= SCALAR_CONVERSION_SHARD_SCHEMA, "compact scalar-conversion admission requires overlay shard schema {}", @@ -556,6 +572,15 @@ pub(super) fn validate_overlay_shard_schema_with_max( "compact TMA reduction admission requires overlay shard schema {}", TMA_REDUCTION_SHARD_SCHEMA ); + ensure!( + shard + .tma + .as_ref() + .is_none_or(|admission| !admission.admits_cta_g2s()) + || shard.schema >= TMA_CTA_SHARD_SCHEMA, + "compact TMA CTA-local G2S admission requires overlay shard schema {}", + TMA_CTA_SHARD_SCHEMA + ); ensure!( shard.mbarrier_extended.is_none() || shard.schema >= MBARRIER_EXTENDED_SHARD_SCHEMA, "compact extended-mbarrier admission requires overlay shard schema {}", 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 082bfbf15e..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 @@ -12,7 +12,8 @@ use crate::model::{ MaskEncoding, MbarrierExtendedAdmission, MbarrierExtendedOperation, MovmatrixAdapter, MovmatrixParticipation, OverlayBackendLowering, OverlayFile, OverlayIntrinsic, PackedAluAdapter, PackedAluFormat, PackedAluOperation, PackedConversionAdapter, - PackedConversionDestinationFormat, PackedConversionFp8Admission, PackedConversionFp8Direction, + PackedConversionDestinationFormat, PackedConversionFp8Admission, + PackedConversionFp8Bf16x2Admission, PackedConversionFp8Direction, PackedConversionFp8F16x2Admission, PackedConversionFp8Format, PackedConversionRounding, PackedConversionSaturation, PackedConversionSourceFormat, PreSm70MemberMaskRule, PrmtAdmission, PrmtMode, RegisterMmaAccumulator, RegisterMmaAmpereFloatAdmission, RegisterMmaF8F6F4Admission, @@ -725,9 +726,9 @@ pub(super) fn packed_conversion_policy( record.memory = "none".into(); record.convergent = false; record.execution_scope = "thread".into(); - let (minimum_ptx, minimum_sm) = packed_conversion_floor(&conversion); + let (minimum_ptx, minimum_sm, _targets) = packed_conversion_floor(&conversion); record.minimum_ptx = minimum_ptx.into(); - record.minimum_sm = Some(minimum_sm.into()); + record.minimum_sm = minimum_sm.map(str::to_owned); record.ptx_result = format!("u{result_width}"); record.ptx_isa_section = "9.7.9.22 Data Movement and Conversion Instructions: cvt".into(); record.ptx_isa_url = "https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt".into(); @@ -740,7 +741,7 @@ pub(super) fn packed_conversion_policy( evidence_profile: "test".into(), targets: None, minimum_ptx: Some(minimum_ptx.into()), - minimum_sm: Some(minimum_sm.into()), + minimum_sm: minimum_sm.map(str::to_owned), }) .collect(); let modifiers = packed_conversion_ptx_modifiers(&conversion); @@ -1410,6 +1411,20 @@ pub(super) fn test_fp8_conversion_admission() -> PackedConversionFp8Admission { } } +pub(super) fn test_fp8_bf16x2_conversion_admission() -> PackedConversionFp8Bf16x2Admission { + PackedConversionFp8Bf16x2Admission { + llvm_evidence_profile: "llvm-test".into(), + libnvvm_evidence_profile: "libnvvm-test".into(), + runtime_validation: RuntimeValidation::Unexecuted, + fp8_formats: vec![ + PackedConversionFp8Format::E4m3x2, + PackedConversionFp8Format::E5m2x2, + ], + relu_variants: true, + product_count: 4, + } +} + pub(super) fn test_fp8_f16x2_conversion_admission() -> PackedConversionFp8F16x2Admission { PackedConversionFp8F16x2Admission { llvm_evidence_profile: "llvm-test".into(), @@ -1471,7 +1486,7 @@ pub(super) fn test_wgmma_control_admission() -> WgmmaControlAdmission { pub(super) fn test_special_register_admission() -> SpecialRegisterAdmission { SpecialRegisterAdmission { - llvm_evidence_profile: "rust-llvm-23.1.0-16696adc".into(), + llvm_evidence_profile: "rust-llvm-23.1.3-aab11ee0".into(), libnvvm_evidence_profile: "cuda-13.3-libnvvm-13.3.33".into(), runtime_validation: RuntimeValidation::Unexecuted, registers: REVIEWED_SPECIAL_REGISTERS.into(), @@ -1520,6 +1535,8 @@ pub(super) fn test_tma_admission() -> TmaAdmission { TmaAdmission { llvm_evidence_profile: "llvm-tma-test".into(), libnvvm_evidence_profile: "libnvvm-tma-test".into(), + cta_llvm_evidence_profile: Some("llvm-tma-test".into()), + cta_libnvvm_evidence_profile: Some("libnvvm-tma-test".into()), reduce_llvm_evidence_profile: Some("llvm-tma-reduce-test".into()), reduce_libnvvm_evidence_profile: Some("libnvvm-tma-reduce-test".into()), runtime_validation: RuntimeValidation::Unexecuted, @@ -1868,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(), @@ -1877,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 8f50c158f2..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 @@ -191,6 +191,7 @@ fn debug_control_compact_schema_is_reserved_for_aggregation() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, @@ -233,7 +234,7 @@ fn active_debug_control_sources_parse_and_prove_both_backend_routes() { let evidence = vec![ read_evidence_file( - &repo_root.join("intrinsics/evidence/rust-llvm-23.1.0-16696adc-debug-control.json"), + &repo_root.join("intrinsics/evidence/rust-llvm-23.1.3-aab11ee0-debug-control.json"), ) .unwrap(), read_evidence_file( @@ -241,7 +242,7 @@ fn active_debug_control_sources_parse_and_prove_both_backend_routes() { ) .unwrap(), ]; - let indexed = index_evidence(&evidence, "16696adcd119e6ba9cc175207d984d7021211acb").unwrap(); + let indexed = index_evidence(&evidence, "aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019").unwrap(); for record in &records { let routes = resolve_backend_lowerings(record, &indexed).unwrap(); assert_eq!(routes.len(), 2); @@ -251,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::>() ); @@ -460,6 +528,7 @@ fn wgmma_control_compact_schema_is_reserved_for_aggregation() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, 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 d62473d0ff..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| { @@ -629,6 +691,18 @@ fn compact_standard_fp8_admission_is_closed_and_ordered() { assert!(expand_register_mma_fp8_admission(&wrong).is_err()); } +/// PTX ISA floor for each Ampere floating-point MMA form; the FP16 m16n8k8 +/// forms date from Turing while the others need Ampere +fn ampere_float_mma_floor(id: &str) -> (&'static str, &'static str) { + match id { + "mma_m16n8k8_f16_f16" | "mma_m16n8k8_f32_f16" => ("6.5", "sm_75"), + "mma_m16n8k4_f32_tf32" | "mma_m16n8k8_f32_bf16" | "mma_m16n8k16_f16_f16" => { + ("7.0", "sm_80") + } + other => panic!("{other} is not an Ampere floating-point MMA form"), + } +} + #[test] fn compact_ampere_float_mma_admission_is_closed_and_ordered() { let admission = test_register_mma_ampere_float_admission(); @@ -653,13 +727,14 @@ fn compact_ampere_float_mma_admission_is_closed_and_ordered() { ] ); for record in &records { - assert_eq!(record.minimum_ptx, "7.0"); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); + let (minimum_ptx, minimum_sm) = ampere_float_mma_floor(&record.id); + assert_eq!(record.minimum_ptx, minimum_ptx); + assert_eq!(record.minimum_sm.as_deref(), Some(minimum_sm)); assert_eq!(record.backend_lowerings.len(), 2); assert!(record.backend_lowerings.iter().all(|lowering| { lowering.mechanism == BackendLoweringMechanism::InlinePtx - && lowering.minimum_ptx.as_deref() == Some("7.0") - && lowering.minimum_sm.as_deref() == Some("sm_80") + && lowering.minimum_ptx.as_deref() == Some(minimum_ptx) + && lowering.minimum_sm.as_deref() == Some(minimum_sm) })); let mma = record.register_mma.as_ref().unwrap(); assert_eq!(mma.kind, None); @@ -734,8 +809,23 @@ fn ampere_float_mma_policies_match_llvm_and_fail_closed() { ], }; assert_eq!(declaration.selections[0].predicates, expected_predicates); - assert_eq!(record.minimum_ptx, "7.0"); - assert_eq!(record.minimum_sm.as_deref(), Some("sm_80")); + let (minimum_ptx, minimum_sm) = ampere_float_mma_floor(&record.id); + assert_eq!(record.minimum_ptx, minimum_ptx); + assert_eq!(record.minimum_sm.as_deref(), Some(minimum_sm)); + } + + // the FP16 m16n8k8 forms reject the old family-wide sm_80 floor + for fp16 in records.iter().filter(|record| { + matches!( + record.id.as_str(), + "mma_m16n8k8_f16_f16" | "mma_m16n8k8_f32_f16" + ) + }) { + let declaration = declarations[fp16.source_record.as_deref().unwrap()]; + let mut wrong = fp16.clone(); + wrong.minimum_ptx = "7.0".into(); + wrong.minimum_sm = Some("sm_80".into()); + assert!(validate_imported_policy(&wrong, declaration).is_err()); } let valid = &records[2]; @@ -757,6 +847,11 @@ fn ampere_float_mma_policies_match_llvm_and_fail_closed() { let mut wrong = valid.clone(); wrong.minimum_ptx = "6.5".into(); assert!(validate_imported_policy(&wrong, declaration).is_err()); + // LLVM selects BF16 at sm_75, but the PTX ISA requires sm_80 + let mut wrong = valid.clone(); + wrong.minimum_ptx = "6.5".into(); + wrong.minimum_sm = Some("sm_75".into()); + assert!(validate_imported_policy(&wrong, declaration).is_err()); let mut wrong = valid.clone(); wrong.register_mma.as_mut().unwrap().adapter = RegisterMmaAdapter::C4F32A4U32B2U32ToD4F32; assert!(validate_imported_policy(&wrong, declaration).is_err()); @@ -969,7 +1064,7 @@ fn standard_fp8_resolves_exact_routes_and_target_floors() { ( IntrinsicBackend::LlvmNvptx, BackendLoweringMechanism::InlinePtx, - "rust-llvm-23.1.0-16696adc", + "rust-llvm-23.1.3-aab11ee0", ), ( IntrinsicBackend::LibNvvm, @@ -1028,7 +1123,11 @@ fn standard_fp8_resolves_exact_routes_and_target_floors() { assert!(llvm_route.stages.iter().any(|stage| { stage.stage == EvidenceStageKind::PtxAssembly && stage.mechanism == Some(BackendLoweringMechanism::InlinePtx) - && stage.outcome == "failed" + && stage.outcome == "not_revalidated" + && stage.detail.contains("Original outcome: failed.") + && stage + .detail + .contains("e457a7b0d326d67b4322ef0d11bd715cfaeda48f") && stage.targets == ["sm_89", "ptx86"] })); } @@ -1097,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 44707537ae..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(), 66); - assert_eq!(overlay.intrinsics.len(), 1029); + assert_eq!(overlay.shards.len(), 68); + assert_eq!(overlay.intrinsics.len(), 1042); assert_eq!( overlay .intrinsics @@ -104,7 +104,7 @@ fn overlay_manifest_loads_sorted_family_shards() { .iter() .filter(|record| record.family == "packed_conversion") .count(), - 18 + 22 ); assert_eq!( overlay @@ -136,7 +136,7 @@ fn overlay_manifest_loads_sorted_family_shards() { .iter() .filter(|record| record.family == "register_mma") .count(), - 154 + 155 ); assert_eq!( overlay @@ -268,6 +268,7 @@ fn overlay_shard_schema_range_is_composable_and_new_fields_fail_closed() { sparse_mma_ordered_ampere_float: None, prmt, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, @@ -343,6 +344,17 @@ fn overlay_shard_schema_range_is_composable_and_new_fields_fail_closed() { .to_string() .contains("requires overlay shard schema 47") ); + let mut cta_tma = shard(TMA_CTA_SHARD_SCHEMA, None, None); + cta_tma.family = "tma".into(); + cta_tma.tma = Some(test_tma_admission()); + validate_overlay_shard_schema(&cta_tma, path).unwrap(); + cta_tma.schema -= 1; + assert!( + validate_overlay_shard_schema(&cta_tma, path) + .unwrap_err() + .to_string() + .contains("compact TMA CTA-local G2S admission requires overlay shard schema 65") + ); let mut sparse_f16_mma = shard(SPARSE_MMA_F8F6F4_F16_SHARD_SCHEMA, None, None); sparse_f16_mma.sparse_mma_f8f6f4_f16 = Some(test_sparse_mma_f8f6f4_f16_admission()); @@ -409,6 +421,7 @@ fn overlay_shard_schema_range_is_composable_and_new_fields_fail_closed() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: Some(test_fp8_conversion_admission()), + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, @@ -455,6 +468,7 @@ fn overlay_shard_schema_range_is_composable_and_new_fields_fail_closed() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, @@ -502,6 +516,7 @@ fn overlay_shard_schema_range_is_composable_and_new_fields_fail_closed() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/scalar_packed.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/scalar_packed.rs index 4d7b0dc448..a1ad55b83b 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/scalar_packed.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/scalar_packed.rs @@ -171,6 +171,71 @@ fn compact_fp8_f16x2_conversion_axes_require_the_exact_closed_product() { assert!(expand_packed_conversion_fp8_f16x2_admission(&executed).is_err()); } +#[test] +fn compact_fp8_bf16x2_pack_axes_require_the_closed_rn_satfinite_product() { + let records = + expand_packed_conversion_fp8_bf16x2_admission(&test_fp8_bf16x2_conversion_admission()) + .unwrap(); + assert_eq!(records.len(), 4); + assert_eq!(records[0].id, "cvt_rn_satfinite_e4m3x2_bf16x2"); + assert_eq!(records[0].rust_name, "cvt_rn_satfinite_e4m3x2_bf16x2"); + assert_eq!( + records[0].expected_ptx.modifiers, + ["rn", "satfinite", "e4m3x2", "bf16x2"] + ); + assert_eq!(records[1].rust_name, "cvt_rn_satfinite_relu_e4m3x2_bf16x2"); + assert_eq!( + records[1].expected_ptx.modifiers, + ["rn", "satfinite", "relu", "e4m3x2", "bf16x2"] + ); + assert_eq!(records[2].rust_name, "cvt_rn_satfinite_e5m2x2_bf16x2"); + assert_eq!( + records[2].expected_ptx.modifiers, + ["rn", "satfinite", "e5m2x2", "bf16x2"] + ); + assert_eq!(records[3].rust_name, "cvt_rn_satfinite_relu_e5m2x2_bf16x2"); + assert_eq!( + records[3].expected_ptx.modifiers, + ["rn", "satfinite", "relu", "e5m2x2", "bf16x2"] + ); + assert!(records.iter().all(|record| { + record.rust_arguments == ["u32"] + && record.rust_result == "u16" + && record.llvm_arguments == ["v2bf16"] + && record.llvm_results == ["i16"] + && record.expected_ptx.mnemonic == "cvt" + && record.expected_ptx.operands.len() == 2 + && record.minimum_ptx == "9.1" + && record.minimum_sm.is_none() + && record.targets + == "sm_100a|sm_100f|sm_103a|sm_103f|sm_110a|sm_110f|sm_120a|sm_120f|sm_121a|sm_121f" + && record + .source_record + .as_deref() + .is_some_and(|source| source.starts_with("int_nvvm_bf16x2_to_")) + })); + + let mut missing_format = test_fp8_bf16x2_conversion_admission(); + missing_format.fp8_formats.pop(); + assert!(expand_packed_conversion_fp8_bf16x2_admission(&missing_format).is_err()); + + let mut reversed_formats = test_fp8_bf16x2_conversion_admission(); + reversed_formats.fp8_formats.reverse(); + assert!(expand_packed_conversion_fp8_bf16x2_admission(&reversed_formats).is_err()); + + let mut without_relu = test_fp8_bf16x2_conversion_admission(); + without_relu.relu_variants = false; + assert!(expand_packed_conversion_fp8_bf16x2_admission(&without_relu).is_err()); + + let mut wrong_count = test_fp8_bf16x2_conversion_admission(); + wrong_count.product_count = 2; + assert!(expand_packed_conversion_fp8_bf16x2_admission(&wrong_count).is_err()); + + let mut executed = test_fp8_bf16x2_conversion_admission(); + executed.runtime_validation = RuntimeValidation::Executed; + assert!(expand_packed_conversion_fp8_bf16x2_admission(&executed).is_err()); +} + #[test] fn packed_atomic_closed_semantics_reject_every_unreviewed_mutation() { let valid = packed_policy("packed_atomic_add_f16x2"); @@ -424,7 +489,7 @@ fn pinned_packed_conversion_records_match_the_closed_recipes() { .iter() .filter(|record| record.family == "packed_conversion") .collect(); - assert_eq!(packed.len(), 18); + assert_eq!(packed.len(), 22); for policy in packed { let source = resolve_policy_source(policy).unwrap(); let declaration = policy diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/sreg_threadfence.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/sreg_threadfence.rs index 593cd62229..959622a23d 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/sreg_threadfence.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/sreg_threadfence.rs @@ -116,6 +116,7 @@ fn special_register_admission_is_closed_and_schema_gated() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, @@ -245,7 +246,7 @@ fn special_register_evidence_validates_both_backend_routes() { let policies = expand_special_register_admission(&test_special_register_admission()).unwrap(); let mut evidence_files = vec![ read_evidence_file( - &repo_root.join("intrinsics/evidence/rust-llvm-23.1.0-16696adc-special-registers.json"), + &repo_root.join("intrinsics/evidence/rust-llvm-23.1.3-aab11ee0-special-registers.json"), ) .unwrap(), read_evidence_file( @@ -253,7 +254,7 @@ fn special_register_evidence_validates_both_backend_routes() { ) .unwrap(), ]; - let llvm_revision = "16696adcd119e6ba9cc175207d984d7021211acb"; + let llvm_revision = "aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019"; let indexed = index_evidence(&evidence_files, llvm_revision).unwrap(); for policy in &policies { for lowering in &policy.backend_lowerings { diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/tcgen05.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/tcgen05.rs index 91aae6815a..48c6f8c1e3 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/tcgen05.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/tcgen05.rs @@ -1188,6 +1188,7 @@ fn tcgen05_compact_schema_is_reserved_for_aggregation() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, diff --git a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/tma_clc_cluster.rs b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/tma_clc_cluster.rs index ceb416cb46..3aa380c48f 100644 --- a/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/tma_clc_cluster.rs +++ b/cuda-oxide/crates/cuda-intrinsics-gen/src/resolve/tests/tma_clc_cluster.rs @@ -116,6 +116,7 @@ fn clc_compact_schema_is_reserved_for_aggregation() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, @@ -190,12 +191,17 @@ fn compact_tma_admission_matches_llvm_and_fails_closed() { .collect::>(), [ ("i0328", "cp_async_bulk_tensor_1d_g2s"), + ("i1034", "cp_async_bulk_tensor_1d_g2s_cta"), ("i0329", "cp_async_bulk_tensor_2d_g2s"), + ("i1035", "cp_async_bulk_tensor_2d_g2s_cta"), ("i0330", "cp_async_bulk_tensor_2d_g2s_multicast"), ("i0331", "cp_async_bulk_tensor_2d_g2s_multicast_cg2"), ("i0332", "cp_async_bulk_tensor_3d_g2s"), + ("i1036", "cp_async_bulk_tensor_3d_g2s_cta"), ("i0333", "cp_async_bulk_tensor_4d_g2s"), + ("i1037", "cp_async_bulk_tensor_4d_g2s_cta"), ("i0334", "cp_async_bulk_tensor_5d_g2s"), + ("i1038", "cp_async_bulk_tensor_5d_g2s_cta"), ("i0335", "cp_async_bulk_tensor_1d_s2g"), ("i0336", "cp_async_bulk_tensor_2d_s2g"), ("i0337", "cp_async_bulk_tensor_3d_s2g"), @@ -277,7 +283,7 @@ fn compact_tma_admission_matches_llvm_and_fails_closed() { } assert_eq!( - parse_hardware_target(&records[3]).unwrap(), + parse_hardware_target(&records[5]).unwrap(), CatalogHardwareTarget::AnyOf { alternatives: vec![ CatalogHardwareAlternative::ExactArchitecture { sm: 100 }, @@ -317,6 +323,47 @@ fn compact_tma_admission_matches_llvm_and_fails_closed() { executed.runtime_validation = RuntimeValidation::Executed; assert!(expand_tma_admission(&executed).is_err()); + let mut cta_without_profiles = test_tma_admission(); + cta_without_profiles.cta_llvm_evidence_profile = None; + assert!( + expand_tma_admission(&cta_without_profiles) + .unwrap_err() + .to_string() + .contains("requires cta_llvm_evidence_profile and cta_libnvvm_evidence_profile") + ); + + let cta_count = TMA_OPERATIONS + .iter() + .filter(|operation| operation.is_cta_g2s()) + .count(); + let mut cluster_only = test_tma_admission(); + cluster_only + .variants + .retain(|variant| !variant.operation.is_cta_g2s()); + cluster_only.cta_llvm_evidence_profile = None; + cluster_only.cta_libnvvm_evidence_profile = None; + let cluster_only_records = expand_tma_admission(&cluster_only).unwrap(); + assert_eq!( + cluster_only_records.len(), + TMA_OPERATIONS.len() - cta_count + tma_reduction_matrix().len() + ); + assert!( + cluster_only_records + .iter() + .all(|record| !record.id.ends_with("_g2s_cta")) + ); + + let mut partial_cta = test_tma_admission(); + partial_cta.variants.retain(|variant| { + !variant.operation.is_cta_g2s() || variant.operation == TmaOperation::G2sCtaTile2d + }); + assert!( + expand_tma_admission(&partial_cta) + .unwrap_err() + .to_string() + .contains("canonical order") + ); + let declaration = declarations[records[0].source_record.as_deref().unwrap()]; let mut wrong_adapter = records[0].clone(); wrong_adapter.tma.as_mut().unwrap().adapter = TmaAdapter::NoOperands; @@ -325,11 +372,18 @@ fn compact_tma_admission_matches_llvm_and_fails_closed() { #[test] fn tma_compact_schema_is_reserved_for_aggregation() { - let shard = |schema: u32, include_reductions: bool| { + let shard = |schema: u32, include_reductions: bool, include_cta: bool| { let mut admission = test_tma_admission(); if !include_reductions { admission.reduce_variants.clear(); } + if !include_cta { + admission + .variants + .retain(|variant| !variant.operation.is_cta_g2s()); + admission.cta_llvm_evidence_profile = None; + admission.cta_libnvvm_evidence_profile = None; + } OverlayShardFile { schema, family: "tma".into(), @@ -349,6 +403,7 @@ fn tma_compact_schema_is_reserved_for_aggregation() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, @@ -371,14 +426,14 @@ fn tma_compact_schema_is_reserved_for_aggregation() { let path = Path::new("intrinsics/overlay/tma.toml"); validate_overlay_shard_schema_with_max( - &shard(TMA_SHARD_SCHEMA, false), + &shard(TMA_SHARD_SCHEMA, false, false), path, TMA_REDUCTION_SHARD_SCHEMA, ) .unwrap(); assert!( validate_overlay_shard_schema_with_max( - &shard(TMA_SHARD_SCHEMA - 1, false), + &shard(TMA_SHARD_SCHEMA - 1, false, false), path, TMA_REDUCTION_SHARD_SCHEMA, ) @@ -388,14 +443,14 @@ fn tma_compact_schema_is_reserved_for_aggregation() { ); validate_overlay_shard_schema_with_max( - &shard(TMA_REDUCTION_SHARD_SCHEMA, true), + &shard(TMA_REDUCTION_SHARD_SCHEMA, true, false), path, TMA_REDUCTION_SHARD_SCHEMA, ) .unwrap(); assert!( validate_overlay_shard_schema_with_max( - &shard(TMA_REDUCTION_SHARD_SCHEMA - 1, true), + &shard(TMA_REDUCTION_SHARD_SCHEMA - 1, true, false), path, TMA_REDUCTION_SHARD_SCHEMA, ) @@ -403,6 +458,23 @@ fn tma_compact_schema_is_reserved_for_aggregation() { .to_string() .contains("compact TMA reduction admission requires overlay shard schema 62") ); + + validate_overlay_shard_schema_with_max( + &shard(TMA_CTA_SHARD_SCHEMA, true, true), + path, + TMA_CTA_SHARD_SCHEMA, + ) + .unwrap(); + assert!( + validate_overlay_shard_schema_with_max( + &shard(TMA_CTA_SHARD_SCHEMA - 1, true, true), + path, + TMA_CTA_SHARD_SCHEMA, + ) + .unwrap_err() + .to_string() + .contains("compact TMA CTA-local G2S admission requires overlay shard schema 65") + ); } #[test] @@ -426,6 +498,7 @@ fn cluster_memory_compact_schema_is_reserved_and_fail_closed() { sparse_mma_ordered_ampere_float: None, prmt: None, packed_conversion_fp8: None, + packed_conversion_fp8_bf16x2: None, packed_conversion_fp8_f16x2: None, scalar_conversion: None, scalar_arithmetic: None, @@ -632,12 +705,12 @@ fn compact_cluster_barrier_admission_and_semantics_fail_closed() { fn cluster_barrier_evidence_validates_both_backend_routes() { let repo_root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); let mut admission = test_cluster_barrier_admission(); - admission.llvm_evidence_profile = "rust-llvm-23.1.0-16696adc".into(); + admission.llvm_evidence_profile = "rust-llvm-23.1.3-aab11ee0".into(); admission.libnvvm_evidence_profile = "cuda-13.3-libnvvm-13.3.33-cluster-barrier".into(); let policies = expand_cluster_barrier_admission(&admission).unwrap(); let evidence_files = vec![ read_evidence_file( - &repo_root.join("intrinsics/evidence/rust-llvm-23.1.0-16696adc-cluster-barrier.json"), + &repo_root.join("intrinsics/evidence/rust-llvm-23.1.3-aab11ee0-cluster-barrier.json"), ) .unwrap(), read_evidence_file( @@ -646,7 +719,7 @@ fn cluster_barrier_evidence_validates_both_backend_routes() { .unwrap(), ]; let indexed = - index_evidence(&evidence_files, "16696adcd119e6ba9cc175207d984d7021211acb").unwrap(); + index_evidence(&evidence_files, "aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019").unwrap(); for policy in &policies { for lowering in &policy.backend_lowerings { @@ -670,7 +743,7 @@ fn cluster_barrier_evidence_validates_both_backend_routes() { } let indexed = index_evidence( &missing_typed_failure, - "16696adcd119e6ba9cc175207d984d7021211acb", + "aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019", ) .unwrap(); let policy = &policies[0]; 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 ce0289c0dd..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Raw ABI functions recognized by cuda-oxide. @@ -1181,6 +1181,22 @@ pub unsafe fn i0328(_arg0: *mut u8, _arg1: *const u8, _arg2: i32, _arg3: *mut u6 ) } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// Catalog ID: `cp_async_bulk_tensor_1d_g2s_cta`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.1d` from `int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_1d`; expects PTX `cp.async.bulk.tensor.1d.shared::cta.global.tile.mbarrier::complete_tx::bytes [
], [
], [
];`. +/// Available on `sm_90+` targets from PTX 8.6. +/// +/// # Safety +/// The destination and barrier must lie in the issuing CTA's own shared memory, and the tensor map must be a live descriptor for this dimensionality. +/// A generic destination or barrier pointer is converted with `cvta.to.shared`, so a cluster-shared pointer from `map_shared_rank` would yield a wrong local offset; copies into another CTA's shared memory use the cluster form. +/// Keep every object alive until the asynchronous copy completes. Only the designated issuing thread may start this transfer. +#[inline(never)] +pub unsafe fn i1034(_arg0: *mut u8, _arg1: *const u8, _arg2: i32, _arg3: *mut u64) -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034` executed outside device compilation" + ) +} + /// Starts a TMA tile copy from shared to global memory. /// /// Catalog ID: `cp_async_bulk_tensor_1d_s2g`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.s2g.tile.1d` from `int_nvvm_cp_async_bulk_tensor_s2g_tile_1d`; expects PTX `cp.async.bulk.tensor.1d.global.shared::cta.tile.bulk_group [
], [
];`. @@ -1217,6 +1233,28 @@ pub unsafe fn i0329( ) } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// Catalog ID: `cp_async_bulk_tensor_2d_g2s_cta`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.2d` from `int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_2d`; expects PTX `cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes [
], [
], [
];`. +/// Available on `sm_90+` targets from PTX 8.6. +/// +/// # Safety +/// The destination and barrier must lie in the issuing CTA's own shared memory, and the tensor map must be a live descriptor for this dimensionality. +/// A generic destination or barrier pointer is converted with `cvta.to.shared`, so a cluster-shared pointer from `map_shared_rank` would yield a wrong local offset; copies into another CTA's shared memory use the cluster form. +/// Keep every object alive until the asynchronous copy completes. Only the designated issuing thread may start this transfer. +#[inline(never)] +pub unsafe fn i1035( + _arg0: *mut u8, + _arg1: *const u8, + _arg2: i32, + _arg3: i32, + _arg4: *mut u64, +) -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035` executed outside device compilation" + ) +} + /// Starts a multicast TMA tile copy from global to cluster shared memory. /// /// Catalog ID: `cp_async_bulk_tensor_2d_g2s_multicast`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d` from `int_nvvm_cp_async_bulk_tensor_g2s_tile_2d`; expects PTX `cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster [
], [
], [
], ;`. @@ -1298,6 +1336,29 @@ pub unsafe fn i0332( ) } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// Catalog ID: `cp_async_bulk_tensor_3d_g2s_cta`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.3d` from `int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_3d`; expects PTX `cp.async.bulk.tensor.3d.shared::cta.global.tile.mbarrier::complete_tx::bytes [
], [
], [
];`. +/// Available on `sm_90+` targets from PTX 8.6. +/// +/// # Safety +/// The destination and barrier must lie in the issuing CTA's own shared memory, and the tensor map must be a live descriptor for this dimensionality. +/// A generic destination or barrier pointer is converted with `cvta.to.shared`, so a cluster-shared pointer from `map_shared_rank` would yield a wrong local offset; copies into another CTA's shared memory use the cluster form. +/// Keep every object alive until the asynchronous copy completes. Only the designated issuing thread may start this transfer. +#[inline(never)] +pub unsafe fn i1036( + _arg0: *mut u8, + _arg1: *const u8, + _arg2: i32, + _arg3: i32, + _arg4: i32, + _arg5: *mut u64, +) -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036` executed outside device compilation" + ) +} + /// Starts a TMA tile copy from shared to global memory. /// /// Catalog ID: `cp_async_bulk_tensor_3d_s2g`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.s2g.tile.3d` from `int_nvvm_cp_async_bulk_tensor_s2g_tile_3d`; expects PTX `cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group [
], [
];`. @@ -1336,6 +1397,30 @@ pub unsafe fn i0333( ) } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// Catalog ID: `cp_async_bulk_tensor_4d_g2s_cta`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.4d` from `int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_4d`; expects PTX `cp.async.bulk.tensor.4d.shared::cta.global.tile.mbarrier::complete_tx::bytes [
], [
], [
];`. +/// Available on `sm_90+` targets from PTX 8.6. +/// +/// # Safety +/// The destination and barrier must lie in the issuing CTA's own shared memory, and the tensor map must be a live descriptor for this dimensionality. +/// A generic destination or barrier pointer is converted with `cvta.to.shared`, so a cluster-shared pointer from `map_shared_rank` would yield a wrong local offset; copies into another CTA's shared memory use the cluster form. +/// Keep every object alive until the asynchronous copy completes. Only the designated issuing thread may start this transfer. +#[inline(never)] +pub unsafe fn i1037( + _arg0: *mut u8, + _arg1: *const u8, + _arg2: i32, + _arg3: i32, + _arg4: i32, + _arg5: i32, + _arg6: *mut u64, +) -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037` executed outside device compilation" + ) +} + /// Starts a TMA tile copy from shared to global memory. /// /// Catalog ID: `cp_async_bulk_tensor_4d_s2g`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.s2g.tile.4d` from `int_nvvm_cp_async_bulk_tensor_s2g_tile_4d`; expects PTX `cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group [
], [
];`. @@ -1383,6 +1468,32 @@ pub unsafe fn i0334( ) } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +/// +/// Catalog ID: `cp_async_bulk_tensor_5d_g2s_cta`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.5d` from `int_nvvm_cp_async_bulk_tensor_g2s_cta_tile_5d`; expects PTX `cp.async.bulk.tensor.5d.shared::cta.global.tile.mbarrier::complete_tx::bytes [
], [
], [
];`. +/// Available on `sm_90+` targets from PTX 8.6. +/// +/// # Safety +/// The destination and barrier must lie in the issuing CTA's own shared memory, and the tensor map must be a live descriptor for this dimensionality. +/// A generic destination or barrier pointer is converted with `cvta.to.shared`, so a cluster-shared pointer from `map_shared_rank` would yield a wrong local offset; copies into another CTA's shared memory use the cluster form. +/// Keep every object alive until the asynchronous copy completes. Only the designated issuing thread may start this transfer. +#[allow(clippy::too_many_arguments)] +#[inline(never)] +pub unsafe fn i1038( + _arg0: *mut u8, + _arg1: *const u8, + _arg2: i32, + _arg3: i32, + _arg4: i32, + _arg5: i32, + _arg6: i32, + _arg7: *mut u64, +) -> () { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038` executed outside device compilation" + ) +} + /// Starts a TMA tile copy from shared to global memory. /// /// Catalog ID: `cp_async_bulk_tensor_5d_s2g`. Source: LLVM `llvm.nvvm.cp.async.bulk.tensor.s2g.tile.5d` from `int_nvvm_cp_async_bulk_tensor_s2g_tile_5d`; expects PTX `cp.async.bulk.tensor.5d.global.shared::cta.tile.bulk_group [
], [
];`. @@ -3002,6 +3113,18 @@ pub fn i0371(_arg0: f32) -> u32 { ) } +/// Converts packed bf16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order. +/// +/// Catalog ID: `cvt_rn_satfinite_e4m3x2_bf16x2`. Source: LLVM `llvm.nvvm.bf16x2.to.e4m3x2.rn.satfinite` from `int_nvvm_bf16x2_to_e4m3x2_rn_satfinite`; expects PTX `cvt.rn.satfinite.e4m3x2.bf16x2 , ;`. +/// Available on `sm_100a or sm_100f or sm_103a or sm_103f or sm_110a or sm_110f or sm_120a or sm_120f or sm_121a or sm_121f` targets from PTX 9.1. +/// Safe because This conversion has no caller obligations. +#[inline(never)] +pub fn i1030(_arg0: u32) -> u16 { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030` executed outside device compilation" + ) +} + /// Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order. /// /// Catalog ID: `cvt_rn_satfinite_e4m3x2_f16x2`. Source: LLVM `llvm.nvvm.f16x2.to.e4m3x2.rn` from `int_nvvm_f16x2_to_e4m3x2_rn`; expects PTX `cvt.rn.satfinite.e4m3x2.f16x2 , ;`. @@ -3026,6 +3149,18 @@ pub fn i0259(_arg0: f32, _arg1: f32) -> u16 { ) } +/// Converts packed bf16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order. +/// +/// Catalog ID: `cvt_rn_satfinite_e5m2x2_bf16x2`. Source: LLVM `llvm.nvvm.bf16x2.to.e5m2x2.rn.satfinite` from `int_nvvm_bf16x2_to_e5m2x2_rn_satfinite`; expects PTX `cvt.rn.satfinite.e5m2x2.bf16x2 , ;`. +/// Available on `sm_100a or sm_100f or sm_103a or sm_103f or sm_110a or sm_110f or sm_120a or sm_120f or sm_121a or sm_121f` targets from PTX 9.1. +/// Safe because This conversion has no caller obligations. +#[inline(never)] +pub fn i1032(_arg0: u32) -> u16 { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032` executed outside device compilation" + ) +} + /// Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order. /// /// Catalog ID: `cvt_rn_satfinite_e5m2x2_f16x2`. Source: LLVM `llvm.nvvm.f16x2.to.e5m2x2.rn` from `int_nvvm_f16x2_to_e5m2x2_rn`; expects PTX `cvt.rn.satfinite.e5m2x2.f16x2 , ;`. @@ -3050,6 +3185,18 @@ pub fn i0261(_arg0: f32, _arg1: f32) -> u16 { ) } +/// Converts packed bf16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order. +/// +/// Catalog ID: `cvt_rn_satfinite_relu_e4m3x2_bf16x2`. Source: LLVM `llvm.nvvm.bf16x2.to.e4m3x2.rn.relu.satfinite` from `int_nvvm_bf16x2_to_e4m3x2_rn_relu_satfinite`; expects PTX `cvt.rn.satfinite.relu.e4m3x2.bf16x2 , ;`. +/// Available on `sm_100a or sm_100f or sm_103a or sm_103f or sm_110a or sm_110f or sm_120a or sm_120f or sm_121a or sm_121f` targets from PTX 9.1. +/// Safe because This conversion has no caller obligations. +#[inline(never)] +pub fn i1031(_arg0: u32) -> u16 { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031` executed outside device compilation" + ) +} + /// Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order. /// /// Catalog ID: `cvt_rn_satfinite_relu_e4m3x2_f16x2`. Source: LLVM `llvm.nvvm.f16x2.to.e4m3x2.rn.relu` from `int_nvvm_f16x2_to_e4m3x2_rn_relu`; expects PTX `cvt.rn.satfinite.relu.e4m3x2.f16x2 , ;`. @@ -3074,6 +3221,18 @@ pub fn i0260(_arg0: f32, _arg1: f32) -> u16 { ) } +/// Converts packed bf16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order. +/// +/// Catalog ID: `cvt_rn_satfinite_relu_e5m2x2_bf16x2`. Source: LLVM `llvm.nvvm.bf16x2.to.e5m2x2.rn.relu.satfinite` from `int_nvvm_bf16x2_to_e5m2x2_rn_relu_satfinite`; expects PTX `cvt.rn.satfinite.relu.e5m2x2.bf16x2 , ;`. +/// Available on `sm_100a or sm_100f or sm_103a or sm_103f or sm_110a or sm_110f or sm_120a or sm_120f or sm_121a or sm_121f` targets from PTX 9.1. +/// Safe because This conversion has no caller obligations. +#[inline(never)] +pub fn i1033(_arg0: u32) -> u16 { + unreachable!( + "generated CUDA intrinsic `cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033` executed outside device compilation" + ) +} + /// Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order. /// /// Catalog ID: `cvt_rn_satfinite_relu_e5m2x2_f16x2`. Source: LLVM `llvm.nvvm.f16x2.to.e5m2x2.rn.relu` from `int_nvvm_f16x2_to_e5m2x2_rn_relu`; expects PTX `cvt.rn.satfinite.relu.e5m2x2.f16x2 , ;`. @@ -3537,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;`. @@ -5093,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 , [
], ;`. @@ -5240,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 , , ;`. @@ -5693,7 +5898,7 @@ pub unsafe fn i0158(_arg0: [i32; 4], _arg1: [u32; 2], _arg2: u32) -> [i32; 4] { ) } -/// Executes the mma_m16n8k16_f16_f16 Ampere floating-point MMA form. +/// Executes the mma_m16n8k16_f16_f16 dense floating-point MMA form. /// /// Catalog ID: `mma_m16n8k16_f16_f16`. Source: LLVM `llvm.nvvm.mma.m16n8k16.row.col.f16.f16` from `int_nvvm_mma_m16n8k16_row_col_f16_f16`; expects PTX `mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 , , , ;`. /// Available on `sm_80+` targets from PTX 7.0. @@ -8204,7 +8409,7 @@ pub unsafe fn i0123(_arg0: [i32; 4], _arg1: [u32; 4], _arg2: [u32; 2]) -> [i32; ) } -/// Executes the mma_m16n8k4_f32_tf32 Ampere floating-point MMA form. +/// Executes the mma_m16n8k4_f32_tf32 dense floating-point MMA form. /// /// Catalog ID: `mma_m16n8k4_f32_tf32`. Source: LLVM `llvm.nvvm.mma.m16n8k4.row.col.tf32` from `int_nvvm_mma_m16n8k4_row_col_tf32`; expects PTX `mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32 , , , ;`. /// Available on `sm_80+` targets from PTX 7.0. @@ -8222,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 , , , ;`. @@ -8374,10 +8610,10 @@ pub unsafe fn i0155(_arg0: [i32; 4], _arg1: [u32; 4], _arg2: [u32; 2]) -> [i32; ) } -/// Executes the mma_m16n8k8_f16_f16 Ampere floating-point MMA form. +/// Executes the mma_m16n8k8_f16_f16 dense floating-point MMA form. /// /// Catalog ID: `mma_m16n8k8_f16_f16`. Source: LLVM `llvm.nvvm.mma.m16n8k8.row.col.f16.f16` from `int_nvvm_mma_m16n8k8_row_col_f16_f16`; expects PTX `mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 , , , ;`. -/// Available on `sm_80+` targets from PTX 7.0. +/// Available on `sm_75+` targets from PTX 6.5. /// /// # Safety /// All 32 warp lanes must execute the same `mma.sync.aligned` instruction with the same qualifiers, and no lane may have exited. @@ -8392,7 +8628,7 @@ pub unsafe fn i0521(_arg0: [u32; 2], _arg1: [u32; 2], _arg2: u32) -> [u32; 2] { ) } -/// Executes the mma_m16n8k8_f32_bf16 Ampere floating-point MMA form. +/// Executes the mma_m16n8k8_f32_bf16 dense floating-point MMA form. /// /// Catalog ID: `mma_m16n8k8_f32_bf16`. Source: LLVM `llvm.nvvm.mma.m16n8k8.row.col.bf16` from `int_nvvm_mma_m16n8k8_row_col_bf16`; expects PTX `mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 , , , ;`. /// Available on `sm_80+` targets from PTX 7.0. @@ -8410,10 +8646,10 @@ pub unsafe fn i0522(_arg0: [f32; 4], _arg1: [u32; 2], _arg2: u32) -> [f32; 4] { ) } -/// Executes the mma_m16n8k8_f32_f16 Ampere floating-point MMA form. +/// Executes the mma_m16n8k8_f32_f16 dense floating-point MMA form. /// /// Catalog ID: `mma_m16n8k8_f32_f16`. Source: LLVM `llvm.nvvm.mma.m16n8k8.row.col.f32.f32` from `int_nvvm_mma_m16n8k8_row_col_f32_f32`; expects PTX `mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 , , , ;`. -/// Available on `sm_80+` targets from PTX 7.0. +/// Available on `sm_75+` targets from PTX 6.5. /// /// # Safety /// All 32 warp lanes must execute the same `mma.sync.aligned` instruction with the same qualifiers, and no lane may have exited. @@ -9248,7 +9484,7 @@ pub unsafe fn i0169( /// `_arg0`, `_arg1`, and `_arg2` hold this lane's 4-register C, 4-register A, and 4-register B fragments; `_arg3` holds its sparse metadata. /// `_arg4` must be the compile-time constant `0`. /// The operation is register-only and is not a memory fence. -/// Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior. +/// Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior. /// See the [PTX sparse MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp). #[must_use] #[inline(never)] @@ -9274,7 +9510,7 @@ pub unsafe fn i1026( /// `_arg0`, `_arg1`, and `_arg2` hold this lane's 4-register C, 4-register A, and 4-register B fragments; `_arg3` holds its sparse metadata. /// `_arg4` must be the compile-time constant `0`. /// The operation is register-only and is not a memory fence. -/// Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior. +/// Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior. /// See the [PTX sparse MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp). #[must_use] #[inline(never)] @@ -9300,7 +9536,7 @@ pub unsafe fn i1027( /// `_arg0`, `_arg1`, and `_arg2` hold this lane's 4-register C, 4-register A, and 4-register B fragments; `_arg3` holds its sparse metadata. /// `_arg4` must be the compile-time constant `0`. /// The operation is register-only and is not a memory fence. -/// Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior. +/// Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior. /// See the [PTX sparse MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp). #[must_use] #[inline(never)] @@ -9326,7 +9562,7 @@ pub unsafe fn i1028( /// `_arg0`, `_arg1`, and `_arg2` hold this lane's 4-register C, 4-register A, and 4-register B fragments; `_arg3` holds its sparse metadata. /// `_arg4` must be the compile-time constant `0`. /// The operation is register-only and is not a memory fence. -/// Every 4-bit metadata group must encode two distinct 2-bit indices; `0x0`, `0x5`, `0xa`, and `0xf` are undefined behavior. +/// Every 4-bit metadata group must be `0x4`, `0x8`, `0x9`, `0xc`, `0xd`, or `0xe`; any other value is undefined behavior. /// See the [PTX sparse MMA fragment layouts](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma-sp). #[must_use] #[inline(never)] diff --git a/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs b/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs index 3ab1989e8d..5a6623e110 100644 --- a/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs +++ b/cuda-oxide/crates/cuda-intrinsics/src/generated/mod.rs @@ -4,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. #[doc(hidden)] #[path = "abi_v1.rs"] @@ -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. @@ -143,6 +146,10 @@ pub mod convert { pub use crate::__cuda_oxide_intrinsic_abi_v1::i0827 as cvt_rn_relu_f16x2_e4m3x2; pub use crate::__cuda_oxide_intrinsic_abi_v1::i0828 as cvt_rn_f16x2_e5m2x2; pub use crate::__cuda_oxide_intrinsic_abi_v1::i0829 as cvt_rn_relu_f16x2_e5m2x2; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1030 as cvt_rn_satfinite_e4m3x2_bf16x2; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1031 as cvt_rn_satfinite_relu_e4m3x2_bf16x2; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1032 as cvt_rn_satfinite_e5m2x2_bf16x2; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1033 as cvt_rn_satfinite_relu_e5m2x2_bf16x2; } /// Generated `debug` intrinsic source API. @@ -679,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. @@ -1094,6 +1102,11 @@ pub mod tma { pub use crate::__cuda_oxide_intrinsic_abi_v1::i0984 as cp_async_bulk_tensor_reduce_xor_im2col_3d; pub use crate::__cuda_oxide_intrinsic_abi_v1::i0985 as cp_async_bulk_tensor_reduce_xor_im2col_4d; pub use crate::__cuda_oxide_intrinsic_abi_v1::i0986 as cp_async_bulk_tensor_reduce_xor_im2col_5d; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1034 as cp_async_bulk_tensor_1d_g2s_cta; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1035 as cp_async_bulk_tensor_2d_g2s_cta; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1036 as cp_async_bulk_tensor_3d_g2s_cta; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1037 as cp_async_bulk_tensor_4d_g2s_cta; + pub use crate::__cuda_oxide_intrinsic_abi_v1::i1038 as cp_async_bulk_tensor_5d_g2s_cta; } /// Generated `warp` intrinsic source API. @@ -1690,6 +1703,14 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cp_async_bulk_tensor_1d_g2s_cta_reexports_abi_i1034() { + let public: unsafe fn(*mut u8, *const u8, i32, *mut u64) -> () = + super::tma::cp_async_bulk_tensor_1d_g2s_cta; + let canonical: unsafe fn(*mut u8, *const u8, i32, *mut u64) -> () = + super::__cuda_oxide_intrinsic_abi_v1::i1034; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cp_async_bulk_tensor_1d_s2g_reexports_abi_i0335() { let public: unsafe fn(*const u8, *const u8, i32) -> () = super::tma::cp_async_bulk_tensor_1d_s2g; @@ -1706,6 +1727,14 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cp_async_bulk_tensor_2d_g2s_cta_reexports_abi_i1035() { + let public: unsafe fn(*mut u8, *const u8, i32, i32, *mut u64) -> () = + super::tma::cp_async_bulk_tensor_2d_g2s_cta; + let canonical: unsafe fn(*mut u8, *const u8, i32, i32, *mut u64) -> () = + super::__cuda_oxide_intrinsic_abi_v1::i1035; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cp_async_bulk_tensor_2d_g2s_multicast_reexports_abi_i0330() { let public: unsafe fn(*mut u8, *const u8, i32, i32, *mut u64, u16) -> () = super::tma::cp_async_bulk_tensor_2d_g2s_multicast; @@ -1738,6 +1767,14 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cp_async_bulk_tensor_3d_g2s_cta_reexports_abi_i1036() { + let public: unsafe fn(*mut u8, *const u8, i32, i32, i32, *mut u64) -> () = + super::tma::cp_async_bulk_tensor_3d_g2s_cta; + let canonical: unsafe fn(*mut u8, *const u8, i32, i32, i32, *mut u64) -> () = + super::__cuda_oxide_intrinsic_abi_v1::i1036; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cp_async_bulk_tensor_3d_s2g_reexports_abi_i0337() { let public: unsafe fn(*const u8, *const u8, i32, i32, i32) -> () = super::tma::cp_async_bulk_tensor_3d_s2g; @@ -1754,6 +1791,14 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cp_async_bulk_tensor_4d_g2s_cta_reexports_abi_i1037() { + let public: unsafe fn(*mut u8, *const u8, i32, i32, i32, i32, *mut u64) -> () = + super::tma::cp_async_bulk_tensor_4d_g2s_cta; + let canonical: unsafe fn(*mut u8, *const u8, i32, i32, i32, i32, *mut u64) -> () = + super::__cuda_oxide_intrinsic_abi_v1::i1037; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cp_async_bulk_tensor_4d_s2g_reexports_abi_i0338() { let public: unsafe fn(*const u8, *const u8, i32, i32, i32, i32) -> () = super::tma::cp_async_bulk_tensor_4d_s2g; @@ -1770,6 +1815,14 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cp_async_bulk_tensor_5d_g2s_cta_reexports_abi_i1038() { + let public: unsafe fn(*mut u8, *const u8, i32, i32, i32, i32, i32, *mut u64) -> () = + super::tma::cp_async_bulk_tensor_5d_g2s_cta; + let canonical: unsafe fn(*mut u8, *const u8, i32, i32, i32, i32, i32, *mut u64) -> () = + super::__cuda_oxide_intrinsic_abi_v1::i1038; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cp_async_bulk_tensor_5d_s2g_reexports_abi_i0339() { let public: unsafe fn(*const u8, *const u8, i32, i32, i32, i32, i32) -> () = super::tma::cp_async_bulk_tensor_5d_s2g; @@ -2465,6 +2518,12 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cvt_rn_satfinite_e4m3x2_bf16x2_reexports_abi_i1030() { + let public: fn(u32) -> u16 = super::convert::cvt_rn_satfinite_e4m3x2_bf16x2; + let canonical: fn(u32) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i1030; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cvt_rn_satfinite_e4m3x2_f16x2_reexports_abi_i0822() { let public: fn(u32) -> u16 = super::convert::cvt_rn_satfinite_e4m3x2_f16x2; let canonical: fn(u32) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i0822; @@ -2477,6 +2536,12 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cvt_rn_satfinite_e5m2x2_bf16x2_reexports_abi_i1032() { + let public: fn(u32) -> u16 = super::convert::cvt_rn_satfinite_e5m2x2_bf16x2; + let canonical: fn(u32) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i1032; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cvt_rn_satfinite_e5m2x2_f16x2_reexports_abi_i0824() { let public: fn(u32) -> u16 = super::convert::cvt_rn_satfinite_e5m2x2_f16x2; let canonical: fn(u32) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i0824; @@ -2489,6 +2554,12 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cvt_rn_satfinite_relu_e4m3x2_bf16x2_reexports_abi_i1031() { + let public: fn(u32) -> u16 = super::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2; + let canonical: fn(u32) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i1031; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cvt_rn_satfinite_relu_e4m3x2_f16x2_reexports_abi_i0823() { let public: fn(u32) -> u16 = super::convert::cvt_rn_satfinite_relu_e4m3x2_f16x2; let canonical: fn(u32) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i0823; @@ -2501,6 +2572,12 @@ mod tests { assert_eq!(public as usize, canonical as usize); } #[test] + fn public_cvt_rn_satfinite_relu_e5m2x2_bf16x2_reexports_abi_i1033() { + let public: fn(u32) -> u16 = super::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2; + let canonical: fn(u32) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i1033; + assert_eq!(public as usize, canonical as usize); + } + #[test] fn public_cvt_rn_satfinite_relu_e5m2x2_f16x2_reexports_abi_i0825() { let public: fn(u32) -> u16 = super::convert::cvt_rn_satfinite_relu_e5m2x2_f16x2; let canonical: fn(u32) -> u16 = super::__cuda_oxide_intrinsic_abi_v1::i0825; @@ -2718,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; @@ -3431,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; @@ -3495,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; @@ -5121,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 a4cc3fbf4c..440290c02a 100644 --- a/cuda-oxide/crates/cuda-macros/Cargo.toml +++ b/cuda-oxide/crates/cuda-macros/Cargo.toml @@ -40,8 +40,12 @@ quote = { workspace = true } syn = { workspace = true } proc-macro2 = { workspace = true } reserved-oxide-symbols = { workspace = true } +# `#[tile_export]` builds its description with the shared vocabulary and embeds +# it as JSON at expansion time; the serialization goes through that crate. +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/README.md b/cuda-oxide/crates/cuda-macros/README.md index a56047c3c8..19b8eda13c 100644 --- a/cuda-oxide/crates/cuda-macros/README.md +++ b/cuda-oxide/crates/cuda-macros/README.md @@ -412,10 +412,13 @@ typed owned async: ## `gpu_printf!` -- Device-Side Printf -Compiles to CUDA's `vprintf` with C vararg promotion rules. Format string must use C-style specifiers. +Compiles to CUDA's `vprintf` with C vararg promotion rules. A placeholder without a type +character prints its argument by type (`%d` or `%lld` for signed integers, `%u` or `%llu` +for unsigned, `%f` for floats, `%p` for raw pointers); a type character such as `{:x}` or +`{:.2f}` keeps the C conversion you name. ```rust -gpu_printf!("thread %d: val = %f\n", tid as i32, val as f64); +gpu_printf!("thread {}: val = {}\n", tid, val); ``` ## `ptx_asm!` -- Inline PTX diff --git a/cuda-oxide/crates/cuda-macros/src/common.rs b/cuda-oxide/crates/cuda-macros/src/common.rs index 5634dd9ec2..86602f5a6b 100644 --- a/cuda-oxide/crates/cuda-macros/src/common.rs +++ b/cuda-oxide/crates/cuda-macros/src/common.rs @@ -82,6 +82,12 @@ pub(crate) fn call_site_ident_avoiding_item(name: &str, item: &ItemFn) -> Ident /// Returns `Some(compile_error)` to be returned from the macro entry point, /// or `None` if the name is safe. pub(crate) fn reject_reserved_name(name: &Ident) -> Option { + reserved_name_error(name).map(|err| err.to_compile_error().into()) +} + +/// [`reject_reserved_name`] as a `syn::Error`, for macros that compose their +/// diagnostics before converting to a token stream. +pub(crate) fn reserved_name_error(name: &Ident) -> Option { let name_str = name.to_string(); if name_str.starts_with(RESERVED_ROOT) { let msg = format!( @@ -90,7 +96,7 @@ pub(crate) fn reject_reserved_name(name: &Ident) -> Option { is reserved for cuda-oxide internal symbol mangling \ (see cuda-oxide/crates/reserved-oxide-symbols)" ); - Some(syn::Error::new(name.span(), msg).to_compile_error().into()) + Some(syn::Error::new(name.span(), msg)) } else { None } diff --git a/cuda-oxide/crates/cuda-macros/src/device.rs b/cuda-oxide/crates/cuda-macros/src/device.rs index bda0e02b18..27d2a4fdc9 100644 --- a/cuda-oxide/crates/cuda-macros/src/device.rs +++ b/cuda-oxide/crates/cuda-macros/src/device.rs @@ -163,7 +163,9 @@ fn generate_device_function(mut input: ItemFn) -> TokenStream { // // - #[inline(never)] on the prefixed function — ensures each monomorphization // appears as a distinct CGU item so the collector can find it. If it were - // inlined, the function would disappear from the CGU. + // inlined, the function would disappear from the CGU. It is a host-side + // boundary only: the backend does not turn it into device `noinline` + // (`device_codegen::is_generic_device_collector_boundary`). // // - The wrapper forwards type parameters via turbofish: // `cuda_oxide_device__add::(a, b)`. diff --git a/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs b/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs index 940bb3c858..4abb6280ee 100644 --- a/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs +++ b/cuda-oxide/crates/cuda-macros/src/launch_attrs.rs @@ -23,6 +23,7 @@ use syn::{ parse::{Parse, ParseStream}, parse_macro_input, parse_quote, punctuated::Punctuated, + spanned::Spanned, visit::{self, Visit}, visit_mut::{self, VisitMut}, }; @@ -472,10 +473,13 @@ pub(crate) struct LoopUnrollAttrVisitor { impl LoopUnrollAttrVisitor { /// If `attrs` contains an outer `#[unroll]` / `#[unroll(N)]`, remove it and - /// return the parsed factor. Returns `None` when no `unroll` attribute is - /// present (leaving `attrs` untouched). Records a parse error and returns - /// `None` if the attribute is malformed. - fn take_unroll_factor(&mut self, attrs: &mut Vec) -> Option { + /// return the parsed factor with the attribute's span. Returns `None` when + /// no `unroll` attribute is present (leaving `attrs` untouched). Records a + /// parse error and returns `None` if the attribute is malformed. + fn take_unroll_factor( + &mut self, + attrs: &mut Vec, + ) -> Option<(ConstU32Expr, proc_macro2::Span)> { let idx = attrs .iter() .position(|attr| attr.path().is_ident("unroll"))?; @@ -506,14 +510,20 @@ impl LoopUnrollAttrVisitor { if factor.literal_value.is_none() { self.const_expressions.push(factor.clone()); } - Some(factor) + Some((factor, attr.span())) } - /// Build the `__unroll_config::()` marker statement. - fn marker_stmt(factor: &ConstU32Expr) -> Stmt { + /// Build the `__unroll_config::()` marker statement. The call + /// carries `span`, the `#[unroll]` attribute's, so diagnostics about the + /// loop point at the attribute. The braces around the factor keep the + /// macro's own span. On the attribute's span rustc treats them as written + /// by the user, and its `unused_braces` lint then reports every literal + /// factor with a fix-it that deletes the attribute. + fn marker_stmt(factor: &ConstU32Expr, span: proc_macro2::Span) -> Stmt { let expr = &factor.expr; - parse_quote! { - cuda_device::thread::__unroll_config::<{ #expr }>(); + let factor_arg = quote! { { #expr } }; + syn::parse_quote_spanned! {span=> + cuda_device::thread::__unroll_config::<#factor_arg>(); } } } @@ -525,18 +535,27 @@ impl VisitMut for LoopUnrollAttrVisitor { // else falls through to the default recursion below. match expr { Expr::ForLoop(for_loop) => { - if let Some(factor) = self.take_unroll_factor(&mut for_loop.attrs) { - for_loop.body.stmts.insert(0, Self::marker_stmt(&factor)); + if let Some((factor, span)) = self.take_unroll_factor(&mut for_loop.attrs) { + for_loop + .body + .stmts + .insert(0, Self::marker_stmt(&factor, span)); } } Expr::While(while_loop) => { - if let Some(factor) = self.take_unroll_factor(&mut while_loop.attrs) { - while_loop.body.stmts.insert(0, Self::marker_stmt(&factor)); + if let Some((factor, span)) = self.take_unroll_factor(&mut while_loop.attrs) { + while_loop + .body + .stmts + .insert(0, Self::marker_stmt(&factor, span)); } } Expr::Loop(loop_expr) => { - if let Some(factor) = self.take_unroll_factor(&mut loop_expr.attrs) { - loop_expr.body.stmts.insert(0, Self::marker_stmt(&factor)); + if let Some((factor, span)) = self.take_unroll_factor(&mut loop_expr.attrs) { + loop_expr + .body + .stmts + .insert(0, Self::marker_stmt(&factor, span)); } } _ => {} diff --git a/cuda-oxide/crates/cuda-macros/src/lib.rs b/cuda-oxide/crates/cuda-macros/src/lib.rs index 08851b5139..aa4ccd2ac5 100644 --- a/cuda-oxide/crates/cuda-macros/src/lib.rs +++ b/cuda-oxide/crates/cuda-macros/src/lib.rs @@ -39,6 +39,7 @@ mod launch; mod launch_attrs; mod printf; mod ptx_asm; +mod tile_export; #[cfg(test)] mod tests; @@ -72,15 +73,49 @@ use syn::parse_macro_input; /// /// | Specifier | Description | Example | /// |-----------|-----------------|------------------------------------------------| -/// | `{}` | Default format | `gpu_printf!("{}", 42)` | +/// | `{}` | By type (below) | `gpu_printf!("{}", 42)` → "42" | /// | `{:x}` | Hex (lower) | `gpu_printf!("{:x}", 255)` → "ff" | /// | `{:X}` | Hex (upper) | `gpu_printf!("{:X}", 255)` → "FF" | /// | `{:#x}` | Hex with prefix | `gpu_printf!("{:#x}", 255)` → "0xff" | /// | `{:o}` | Octal | `gpu_printf!("{:o}", 8)` → "10" | +/// | `{:d}` | Signed decimal | `gpu_printf!("{:d}", 42u8)` → "42" | /// | `{:e}` | Scientific | `gpu_printf!("{:e}", 1000.0)` → "1.000000e+03" | +/// | `{:c}` | Character | `gpu_printf!("{:c}", 'a')` → "a" | /// | `{:.N}` | Precision | `gpu_printf!("{:.2}", 3.14159)` → "3.14" | /// | `{:N}` | Width | `gpu_printf!("{:8}", 42)` → " 42" | /// | `{:0N}` | Zero-pad | `gpu_printf!("{:08}", 42)` → "00000042" | +/// | `{:*}` | Runtime width | `gpu_printf!("{:*}", 8, 42)` → " 42" | +/// +/// A type character (`x`, `X`, `o`, `d`, `e`, `E`, `f`, `g`, `G`, `a`, `A`, +/// `c`, `p`) picks the C conversion and casts the argument to the matching C +/// type (`as u64`, `as i64`, `as f64`, `as i32`). +/// +/// # Default conversions +/// +/// A placeholder without a type character (`{}`, `{:8}`, `{:08}`, `{:.2}`) +/// prints its argument by type. The argument has to implement +/// `cuda_device::debug::GpuPrintfArg`: +/// +/// | Argument type | Conversion | +/// |----------------------------|------------| +/// | `i8`, `i16`, `i32`, `bool` | `%d` | +/// | `i64`, `isize` | `%lld` | +/// | `u8`, `u16`, `u32` | `%u` | +/// | `u64`, `usize` | `%llu` | +/// | `f32`, `f64` | `%f` | +/// | `*const T`, `*mut T` | `%p` | +/// +/// An unsuffixed integer literal is an `i32` and an unsuffixed float literal +/// an `f64`, as in the rest of Rust. Width, zero padding and precision keep +/// their C meaning: `{:.2}` on a float prints two decimals, and on an integer +/// it is the minimum number of digits (`{:.3}` on `-42` prints `-042`). The +/// `+` and space flags need a signed integer or floating-point argument; on +/// an unsigned or pointer argument they are a compile-time error. +/// +/// Not accepted by a default placeholder: `char` (use `{:c}`), fieldless +/// enums (cast with `as i32`), `i128` and `u128`. Runtime precision `{:.*}` +/// is a compile-time error, because CUDA device printf takes `*` for the +/// width only. /// /// # Returns /// @@ -422,8 +457,9 @@ pub fn cuda_module(attr: TokenStream, item: TokenStream) -> TokenStream { /// invalid specialization fails compilation instead of becoming a no-op. /// /// The pass currently recognizes explicit counted `while` loops and range `for` -/// loops (`for i in a..b`). Other iterator loops, such as `a..=b` or -/// `.step_by(k)`, are not yet recognized. +/// loops (`for i in a..b` and `(a..b).rev()`) over any integer counter, +/// including `isize`. Other iterator loops, such as `a..=b` or `.step_by(k)`, +/// are not yet recognized. /// /// Only the annotated loop is unrolled. Inner loops are copied intact unless /// they carry their own annotation. Several `continue` paths (multiple @@ -837,7 +873,8 @@ pub fn cooperative_launch(attr: TokenStream, item: TokenStream) -> TokenStream { /// /// Loop annotations work the same way in device function definitions as they do /// in kernels. Use an explicit counted `while` loop or a range `for` loop -/// (`for i in a..b`); other iterator loops are not yet recognized. Partial +/// (`for i in a..b` or `(a..b).rev()`, over any integer counter including +/// `isize`); other iterator loops are not yet recognized. Partial /// factors must be `N >= 2`. Multiple `continue` paths are supported; full /// unrolling preserves `break` and multiple exit targets. Partial unrolling /// requires a positive counter step, a `<` or `<=` test, an unchanging limit, @@ -888,6 +925,110 @@ pub fn device(_attr: TokenStream, item: TokenStream) -> TokenStream { device::device_entry(_attr, item) } +/// Exports a device function to cuTile programs through a tile library. +/// +/// A cuTile kernel calls the function with `foreign_call`, naming the +/// declaration; `cargo oxide tilelib` packages a crate's exports into the +/// `.tilelib` that call resolves against. The attribute leaves the annotated +/// function unchanged, so other device code keeps calling it, and generates +/// two items next to it: +/// +/// 1. `#[device] pub unsafe extern "C" fn (...)`: the implementation +/// the Tile IR call resolves to, on the native SIMT ABI below. It turns the +/// pointer parameters back into the references the function takes, calls +/// the function, and stores the return value through the trailing out +/// pointer. +/// 2. A host-only `#[no_mangle]` metadata function that returns the export's +/// JSON description. The codegen backend reads the string out of its MIR, +/// compiles the wrapper as a device root even when the crate also has +/// kernels, and writes `.tileexports.json` next to the device +/// artifact. The metadata function itself is never compiled for the device. +/// +/// # Arguments +/// +/// | Argument | Meaning | Default | +/// |-----------------------------|-------------------------------------------------|-------------------| +/// | `tile_shape = N` | 1-D tile extent (required) | none | +/// | `name = "..."` | declaration name the cuTile caller uses | the function name | +/// | `symbol = "..."` | implementation symbol in the device artifact | `_impl` | +/// | `sync = "none"` / `"warp"` | threads that must call the function together | `"none"` | +/// | `arch = ["sm_120", ...]` | compute capabilities the implementation supports| every architecture| +/// +/// `tile_shape` is required because this version ships concrete implementations: +/// a parametric one would need one mangled symbol per specialization in the +/// device artifact. `sync = "warp"` needs a tile shape that is a multiple of 32 +/// times the parameter footprint, so every invocation group fills whole +/// warps. `arch` is normally omitted, which matches every target; when given, +/// names are compared after stripping `_` and lowercasing, and a portable name +/// such as `"sm_120"` also matches the arch-conditional `sm_120a` the Tile +/// compiler builds for. +/// +/// # Signature +/// +/// The function cannot be generic or take `self`, and every parameter is a +/// plain identifier with one of these types, where `T` is one of `f16`, +/// `f32`, `f64`, `i8`, `i16`, `i32`, `i64`, `u8`, `u16`, `u32`, `u64`: +/// +/// | Rust parameter | access | footprint | ABI parameter | +/// |------------------|---------|-----------|---------------| +/// | `x: T` | `in` | 1 | `T` by value | +/// | `x: &T` | `in` | 1 | `*const T` | +/// | `x: &[T; N]` | `in` | N | `*const T` | +/// | `x: &mut T` | `inout` | 1 | `*mut T` | +/// | `x: &mut [T; N]` | `inout` | N | `*mut T` | +/// +/// The footprint is the number of contiguous tile elements one invocation +/// owns, and every parameter of one export has the same footprint: the +/// fragmaps of one implementation share their invocation count, so +/// `fn window(acc: &mut [i32; 4], bias: &i32)` is rejected and +/// `bias: &[i32; 4]` is the accepted form. +/// +/// A scalar return type `-> T` becomes an appended `out` parameter named +/// `_out`, passed as `*mut T`. Returning nothing is accepted only when some +/// parameter is `&mut`, because a declaration's results are its `out` and +/// `inout` arguments. Slices, raw pointers, tuples, structs, and `bool` are +/// not accepted. `bf16` is not an accepted element type either: `cuda-device` +/// carries bf16 values as `u16` bit patterns, so pass such tiles as `u16`. +/// +/// Thread identity inside the function is the hardware `threadIdx.x`. The +/// function may not allocate shared memory; a scratch buffer is not part of +/// this version of the description. +/// +/// # Example +/// +/// ```ignore +/// use cuda_device::tile_export; +/// +/// #[tile_export(tile_shape = 128)] +/// pub fn relu(x: f32) -> f32 { +/// if x > 0.0 { x } else { 0.0 } +/// } +/// +/// #[tile_export(tile_shape = 128, name = "scale_in_place", symbol = "scale_in_place_impl")] +/// pub fn scale(x: &mut f32, factor: f32) { +/// *x *= factor; +/// } +/// +/// #[tile_export(tile_shape = 128, sync = "warp", arch = ["sm_120"])] +/// pub fn warp_sum(x: f32) -> f32 { +/// // 32 lanes call this together, so shuffles are legal here. +/// let mut total = x; +/// let mut offset = 16; +/// while offset > 0 { +/// total += cuda_device::warp::shuffle_xor_f32_sync(u32::MAX, total, offset); +/// offset /= 2; +/// } +/// total +/// } +/// ``` +/// +/// `relu` becomes `define void @relu_impl(float, ptr)` and `scale` becomes +/// `define void @scale_in_place_impl(ptr, float)` in the device artifact. +#[proc_macro_attribute] +pub fn tile_export(attr: TokenStream, item: TokenStream) -> TokenStream { + tile_export::tile_export_entry(attr, item) +} + // ============================================================================ // NVVM Attributes for Device FFI // ============================================================================ diff --git a/cuda-oxide/crates/cuda-macros/src/printf.rs b/cuda-oxide/crates/cuda-macros/src/printf.rs index 26043dc0d7..dca61461e2 100644 --- a/cuda-oxide/crates/cuda-macros/src/printf.rs +++ b/cuda-oxide/crates/cuda-macros/src/printf.rs @@ -11,30 +11,57 @@ //! # Design //! //! The `gpu_printf!` macro: -//! 1. Parses Rust-style format string (`{}`, `{:x}`, `{:.2}`, etc.) -//! 2. Translates to C format string (`%d`, `%x`, `%.2f`, etc.) -//! 3. Generates a `#[repr(C)]` struct for argument packing -//! 4. Calls `__gpu_vprintf(format_ptr, args_ptr)` which the compiler recognizes +//! 1. Parses the Rust-style format string (`{}`, `{:x}`, `{:.2}`, and so on) +//! 2. Evaluates every argument, in order, into a local binding +//! 3. Packs the bindings into a `#[repr(C)]` struct +//! 4. Calls `__gpu_vprintf(format_ptr, args_ptr)`, which the compiler lowers +//! to `vprintf` //! -//! # C Vararg Promotion Rules +//! # Choosing the C conversion +//! +//! A placeholder with a type character (`{:x}`, `{:.2f}`, `{:d}`, `{:c}`, ...) +//! has a fixed conversion. The macro casts the argument to the matching C +//! vararg type (`as u64`, `as f64`, `as i64`, `as i32`) and, when every +//! placeholder is of this kind, writes the whole format string as a byte +//! string literal. +//! +//! A placeholder without a type character (`{}`, `{:8}`, `{:08}`, `{:.2}`) +//! takes the conversion from the argument's type through +//! `cuda_device::debug::GpuPrintfArg`: `%d` or `%lld` for signed integers, +//! `%u` or `%llu` for unsigned ones, `%f` for floats, `%d` for `bool` and `%p` +//! for raw pointers. The macro cannot see that type, so it emits a generic +//! helper with one `GpuPrintfArg` parameter per such placeholder. An inline +//! const inside the helper calls `cuda_device::debug::__gpu_printf_format` +//! with the literal text and the per-placeholder specs as data, and the +//! trait's constants pick the type character after monomorphization. The +//! helper returns `&'static [u8; N]`, the pointer-to-array constant the device +//! importer accepts. `N` is computed here from the same data and leaves room +//! for an `ll` modifier on every typed placeholder. +//! +//! # C vararg promotion rules //! //! Arguments are promoted following C vararg rules: //! - `i8`, `i16` → `i32` //! - `u8`, `u16` → `u32` -//! - `f32` → `f64` (critical!) +//! - `f32` → `f64` //! - `bool` → `i32` (0 or 1) -//! - pointers/references → `u64` (64-bit) +//! - pointers → `u64` -use proc_macro2::TokenStream as TokenStream2; +use proc_macro2::{Span, TokenStream as TokenStream2}; use quote::{format_ident, quote}; use syn::{ - Expr, LitStr, Token, + Expr, LitByteStr, LitStr, Token, parse::{Parse, ParseStream}, }; /// Maximum number of printf arguments (CUDA limitation) const MAX_PRINTF_ARGS: usize = 32; +/// Bytes a typed conversion needs beyond its prefix: `%`, `ll` and the type +/// character. The builder writes `ll` only for 64-bit integers; the rest of +/// the room stays zero after the NUL. +const TYPED_SPEC_OVERHEAD: usize = 4; + /// Input to gpu_printf! macro: format string followed by arguments pub struct GpuPrintfInput { format_string: LitStr, @@ -86,8 +113,9 @@ struct FormatSpec { zero_pad: bool, /// Minimum width (None = no width specified) width: Option, - /// Precision (None = no precision specified) - precision: Option, + /// Literal precision (None = no precision specified). Device printf + /// accepts digits only here, so there is no runtime form. + precision: Option, /// Format type character (x, X, o, e, E, f, g, G, a, A, p, c, s, d, etc.) format_type: Option, } @@ -101,15 +129,6 @@ enum Width { Runtime, } -/// Precision specification -#[derive(Debug, Clone)] -enum Precision { - /// Literal precision value - Literal(usize), - /// Runtime precision (`.*`) - Runtime, -} - /// Parse all placeholders from a format string fn parse_format_string(format: &str) -> Result, String> { let mut placeholders = Vec::new(); @@ -209,21 +228,24 @@ fn parse_format_spec(chars: &[char], mut i: usize, spec: &mut FormatSpec) -> Res } } - // Parse precision + // Parse precision. CUDA device printf takes `*` for the width only; a + // runtime precision prints its own specifier text instead of a value. if i < chars.len() && chars[i] == '.' { i += 1; // Skip `.` - if i < chars.len() { - if chars[i] == '*' { - spec.precision = Some(Precision::Runtime); - i += 1; - } else if chars[i].is_ascii_digit() { - let (prec, new_i) = parse_number(chars, i); - spec.precision = Some(Precision::Literal(prec)); - i = new_i; - } else { - // `.` without number means precision 0 - spec.precision = Some(Precision::Literal(0)); - } + if i < chars.len() && chars[i] == '*' { + return Err( + "runtime precision `.*` is not supported by CUDA device printf; \ + use a literal precision such as `{:.2}`" + .to_string(), + ); + } + if i < chars.len() && chars[i].is_ascii_digit() { + let (prec, new_i) = parse_number(chars, i); + spec.precision = Some(prec); + i = new_i; + } else { + // `.` without number means precision 0 + spec.precision = Some(0); } } @@ -266,163 +288,115 @@ fn parse_number(chars: &[char], mut i: usize) -> (usize, usize) { (num, i) } -/// Count how many extra arguments are needed for runtime width/precision -fn count_runtime_args(spec: &FormatSpec) -> usize { - let mut count = 0; - if matches!(spec.width, Some(Width::Runtime)) { - count += 1; - } - if matches!(spec.precision, Some(Precision::Runtime)) { - count += 1; +/// Literal text around the placeholders, with `%` doubled and `{{` / `}}` +/// unescaped. Returns one more part than there are placeholders. +fn literal_parts(format: &str, placeholders: &[Placeholder]) -> Vec { + let chars: Vec = format.chars().collect(); + let mut parts = Vec::with_capacity(placeholders.len() + 1); + let mut last_end = 0; + for placeholder in placeholders { + parts.push(unescape_literal(&chars[last_end..placeholder.start])); + last_end = placeholder.end; } - count + parts.push(unescape_literal(&chars[last_end..])); + parts } -/// Build the C format string from Rust format string -fn build_c_format_string(format: &str, placeholders: &[Placeholder]) -> String { +/// One run of literal text, converted to C printf text. +fn unescape_literal(chars: &[char]) -> String { let mut result = String::new(); - let chars: Vec = format.chars().collect(); - let mut last_end = 0; - - for placeholder in placeholders { - // Add literal text before this placeholder - let mut i = last_end; - while i < placeholder.start { - if chars[i] == '{' && i + 1 < chars.len() && chars[i + 1] == '{' { + let mut i = 0; + while i < chars.len() { + match chars[i] { + '{' if i + 1 < chars.len() && chars[i + 1] == '{' => { result.push('{'); i += 2; - } else if chars[i] == '}' && i + 1 < chars.len() && chars[i + 1] == '}' { + } + '}' if i + 1 < chars.len() && chars[i + 1] == '}' => { result.push('}'); i += 2; - } else if chars[i] == '%' { + } + '%' => { // Escape literal % as %% result.push_str("%%"); i += 1; - } else { - result.push(chars[i]); + } + c => { + result.push(c); i += 1; } } - - // Add placeholder as %... (will be filled in by generate_c_specifier) - // For now, add a marker that we'll replace - result.push_str("__PLACEHOLDER__"); - - last_end = placeholder.end; - } - - // Add remaining literal text - let mut i = last_end; - while i < chars.len() { - if chars[i] == '{' && i + 1 < chars.len() && chars[i + 1] == '{' { - result.push('{'); - i += 2; - } else if chars[i] == '}' && i + 1 < chars.len() && chars[i + 1] == '}' { - result.push('}'); - i += 2; - } else if chars[i] == '%' { - result.push_str("%%"); - i += 1; - } else { - result.push(chars[i]); - i += 1; - } } - result } -/// Generate C format specifier from FormatSpec. -fn generate_c_specifier( - spec: &FormatSpec, - is_64bit: bool, - is_float: bool, - is_string: bool, -) -> String { - let mut result = String::from("%"); - - // Add flags +/// Flags, width and precision of a placeholder, as they appear between the +/// `%` and the type character. +fn spec_prefix(spec: &FormatSpec) -> String { + let mut prefix = String::new(); if spec.left_justify { - result.push('-'); + prefix.push('-'); } if spec.always_sign { - result.push('+'); + prefix.push('+'); } if spec.space_sign { - result.push(' '); + prefix.push(' '); } if spec.alternate { - result.push('#'); + prefix.push('#'); } if spec.zero_pad { - result.push('0'); + prefix.push('0'); } - - // Add width match &spec.width { - Some(Width::Literal(w)) => result.push_str(&w.to_string()), - Some(Width::Runtime) => result.push('*'), + Some(Width::Literal(w)) => prefix.push_str(&w.to_string()), + Some(Width::Runtime) => prefix.push('*'), None => {} } - - // Add precision - match &spec.precision { - Some(Precision::Literal(p)) => { - result.push('.'); - result.push_str(&p.to_string()); - } - Some(Precision::Runtime) => result.push_str(".*"), - None => {} + if let Some(p) = spec.precision { + prefix.push('.'); + prefix.push_str(&p.to_string()); } + prefix +} + +/// The complete C conversion of a placeholder with an explicit type character. +fn fixed_specifier( + spec: &FormatSpec, + type_char: char, + is_64bit: bool, + is_float: bool, + is_string: bool, +) -> String { + let mut result = String::from("%"); + result.push_str(&spec_prefix(spec)); // Add size modifier for 64-bit integers if is_64bit && !is_float && !is_string { result.push_str("ll"); } - // Add type character - let type_char = match spec.format_type { - Some('c') => 'c', - Some('s') => 's', - Some('d') | Some('i') => 'd', - Some('o') => 'o', - Some('x') => 'x', - Some('X') => 'X', - Some('e') => 'e', - Some('E') => 'E', - Some('f') => 'f', - Some('g') => 'g', - Some('G') => 'G', - Some('a') => 'a', - Some('A') => 'A', - Some('p') => 'p', - Some('b') => 'b', // Binary - we'll handle this specially - Some(_) => { - // Unknown format specifier - fall back to default - if is_string { - 's' - } else if is_float { - 'f' - } else { - 'd' - } - } - None => { - // Default type based on argument type - if is_string { - 's' - } else if is_float { - 'f' - } else { - 'd' // Will be 'u' for unsigned, handled at arg level - } - } - }; - result.push(type_char); - + result.push(if type_char == 'i' { 'd' } else { type_char }); result } +/// A NUL-terminated byte string literal: the form of a static format string +/// the device accepts (a `str` would not be). +fn c_string_literal(text: &str) -> LitByteStr { + LitByteStr::new(format!("{text}\0").as_bytes(), Span::call_site()) +} + +/// How one placeholder prints its argument. +enum Conversion { + /// Explicit type character: the macro knows the whole C conversion. + Fixed(String), + /// No type character: the argument's `GpuPrintfArg` impl supplies the + /// type character after monomorphization. `prefix` is the text between + /// the `%` and that character, `ty` the helper's type parameter. + Typed { prefix: String, ty: syn::Ident }, +} + /// Main proc macro implementation (called from lib.rs) pub fn gpu_printf_impl(input: GpuPrintfInput) -> TokenStream2 { let format_str = input.format_string.value(); @@ -435,12 +409,11 @@ pub fn gpu_printf_impl(input: GpuPrintfInput) -> TokenStream2 { } }; - // Count total expected arguments (including runtime width/precision) - let mut expected_args = 0; - for p in &placeholders { - expected_args += 1; // The value itself - expected_args += count_runtime_args(&p.spec); - } + // Count total expected arguments (including runtime width) + let expected_args: usize = placeholders + .iter() + .map(|p| 1 + usize::from(matches!(p.spec.width, Some(Width::Runtime)))) + .sum(); // Validate argument count if input.args.len() != expected_args { @@ -468,18 +441,14 @@ pub fn gpu_printf_impl(input: GpuPrintfInput) -> TokenStream2 { .to_compile_error(); } + let parts = literal_parts(&format_str, &placeholders); + // If no arguments, simple case if placeholders.is_empty() { - let c_format = build_c_format_string(&format_str, &placeholders); - // Create byte string literal for format string (avoids str type) - let format_with_null = format!("{}\0", c_format); - let format_lit = - syn::LitByteStr::new(format_with_null.as_bytes(), proc_macro2::Span::call_site()); - + let format_lit = c_string_literal(&parts[0]); return quote! { { // No arguments - pass null for args pointer - // Use byte string literal directly (avoids str type which GPU doesn't support) cuda_device::debug::__gpu_vprintf( #format_lit.as_ptr(), core::ptr::null() @@ -488,10 +457,16 @@ pub fn gpu_printf_impl(input: GpuPrintfInput) -> TokenStream2 { }; } - // Generate the argument struct fields and initializers + // Every argument is evaluated into a local, in source order, before the + // packing struct is built. The locals carry mixed-site spans so a later + // argument expression cannot see an earlier binding. + let mut hoisted: Vec = Vec::new(); let mut field_defs: Vec = Vec::new(); let mut field_inits: Vec = Vec::new(); - let mut c_specifiers: Vec = Vec::new(); + let mut conversions: Vec = Vec::new(); + let mut typed_params: Vec = Vec::new(); + let mut promoted_params: Vec = Vec::new(); + let mut typed_values: Vec = Vec::new(); let mut arg_idx = 0; for (placeholder_idx, placeholder) in placeholders.iter().enumerate() { @@ -499,129 +474,185 @@ pub fn gpu_printf_impl(input: GpuPrintfInput) -> TokenStream2 { // Handle runtime width if matches!(spec.width, Some(Width::Runtime)) { - let field_name = format_ident!("__width_{}", placeholder_idx); + let value = format_ident!("__val_{}", arg_idx, span = Span::mixed_site()); + let field = format_ident!("__width_{}", placeholder_idx); let arg = &input.args[arg_idx]; - field_defs.push(quote! { #field_name: i32 }); - field_inits.push(quote! { #field_name: (#arg) as i32 }); - arg_idx += 1; - } - - // Handle runtime precision - if matches!(spec.precision, Some(Precision::Runtime)) { - let field_name = format_ident!("__prec_{}", placeholder_idx); - let arg = &input.args[arg_idx]; - field_defs.push(quote! { #field_name: i32 }); - field_inits.push(quote! { #field_name: (#arg) as i32 }); + hoisted.push(quote! { let #value = (#arg) as i32; }); + field_defs.push(quote! { #field: i32 }); + field_inits.push(quote! { #field: #value }); arg_idx += 1; } // Handle the value argument - let field_name = format_ident!("__arg_{}", placeholder_idx); + let value = format_ident!("__val_{}", arg_idx, span = Span::mixed_site()); + let field = format_ident!("__arg_{}", placeholder_idx); let arg = &input.args[arg_idx]; arg_idx += 1; - // Generate field type and initialization based on format specifier - // Use concrete types - C vararg promotion rules apply: - // - f32 → f64 (always) - // - i8, i16 → i32; u8, u16 → u32 - // - i32, u32, i64, u64 → unchanged - // We use i64/u64 as the default to handle all integer sizes safely - let (field_ty, field_init, is_64bit, is_float, is_string) = match spec.format_type { - Some('c') => { - // Character: promote to i32 - ( - quote! { i32 }, - quote! { (#arg) as i32 }, - false, - false, - false, - ) - } - Some('s') => { - // String: pass pointer - ( - quote! { *const u8 }, - quote! { { - // Ensure string is null-terminated - const __STR: &str = concat!(#arg, "\0"); - __STR.as_ptr() - }}, - false, - false, - true, - ) - } - Some('p') => { - // Pointer - (quote! { u64 }, quote! { (#arg) as u64 }, true, false, false) - } - Some('e') | Some('E') | Some('f') | Some('g') | Some('G') | Some('a') | Some('A') => { - // Float: always promote to f64 - (quote! { f64 }, quote! { (#arg) as f64 }, false, true, false) - } - Some('x') | Some('X') | Some('o') | Some('b') => { - // Hex/octal/binary: use u64 to handle all unsigned integer sizes - (quote! { u64 }, quote! { (#arg) as u64 }, true, false, false) - } - Some('d') | Some('i') => { - // Signed integer: use i64 to handle all signed integer sizes - (quote! { i64 }, quote! { (#arg) as i64 }, true, false, false) - } + match spec.format_type { None => { - // Default format - use i64 for integers (most common case) - // The C format string will be %lld - (quote! { i64 }, quote! { (#arg) as i64 }, true, false, false) + let ty = format_ident!("__T{}", typed_params.len()); + let promoted = format_ident!("__P{}", typed_params.len()); + hoisted.push(quote! { let #value = #arg; }); + field_defs.push(quote! { #field: #promoted }); + field_inits.push(quote! { + #field: cuda_device::debug::GpuPrintfArg::promote(#value) + }); + typed_params.push(ty.clone()); + promoted_params.push(promoted); + typed_values.push(value); + conversions.push(Conversion::Typed { + prefix: spec_prefix(spec), + ty, + }); } - _ => { - // Fallback to i64 - (quote! { i64 }, quote! { (#arg) as i64 }, true, false, false) + Some(type_char) => { + // Explicit type character: cast to the C vararg type. + // - f32 → f64 (always) + // - integers → i64 or u64, so every width fits + let (field_ty, init, is_64bit, is_float, is_string) = match type_char { + 'c' => ( + quote! { i32 }, + quote! { (#arg) as i32 }, + false, + false, + false, + ), + 's' => ( + quote! { *const u8 }, + quote! { { + // Ensure string is null-terminated + const __STR: &str = concat!(#arg, "\0"); + __STR.as_ptr() + } }, + false, + false, + true, + ), + 'p' => (quote! { u64 }, quote! { (#arg) as u64 }, true, false, false), + 'e' | 'E' | 'f' | 'g' | 'G' | 'a' | 'A' => { + (quote! { f64 }, quote! { (#arg) as f64 }, false, true, false) + } + 'x' | 'X' | 'o' | 'b' => { + (quote! { u64 }, quote! { (#arg) as u64 }, true, false, false) + } + // `d`, `i` and anything else the parser accepts + _ => (quote! { i64 }, quote! { (#arg) as i64 }, true, false, false), + }; + hoisted.push(quote! { let #value = #init; }); + field_defs.push(quote! { #field: #field_ty }); + field_inits.push(quote! { #field: #value }); + conversions.push(Conversion::Fixed(fixed_specifier( + spec, type_char, is_64bit, is_float, is_string, + ))); } - }; + } + } - field_defs.push(quote! { #field_name: #field_ty }); - field_inits.push(quote! { #field_name: #field_init }); + let args_ident = format_ident!("__args", span = Span::mixed_site()); - // Generate C format specifier - // For trait-based types, we use a placeholder that gets resolved at compile time - let c_spec = generate_c_specifier(spec, is_64bit, is_float, is_string); - c_specifiers.push(c_spec); - } + if typed_params.is_empty() { + // Every conversion is known here: emit the format string as a literal. + let mut c_format = String::new(); + for (part, conversion) in parts.iter().zip(&conversions) { + c_format.push_str(part); + match conversion { + Conversion::Fixed(spec) => c_format.push_str(spec), + Conversion::Typed { .. } => unreachable!("no typed placeholders on this route"), + } + } + c_format.push_str(&parts[conversions.len()]); + let format_lit = c_string_literal(&c_format); - // Build the final C format string - let mut c_format = build_c_format_string(&format_str, &placeholders); - for c_spec in &c_specifiers { - c_format = c_format.replacen("__PLACEHOLDER__", c_spec, 1); - } + return quote! { + { + // Argument packing struct with C-compatible layout + #[repr(C)] + struct __GpuPrintfArgs { + #(#field_defs),* + } + + #(#hoisted)* + + let #args_ident = __GpuPrintfArgs { + #(#field_inits),* + }; - // Generate unique struct name - let struct_name = format_ident!("__GpuPrintfArgs"); + cuda_device::debug::__gpu_vprintf( + #format_lit.as_ptr(), + &#args_ident as *const __GpuPrintfArgs as *const u8 + ) + } + }; + } - // Create byte string literal for format string (avoids str type) - let format_with_null = format!("{}\0", c_format); - let format_lit = - syn::LitByteStr::new(format_with_null.as_bytes(), proc_macro2::Span::call_site()); + // At least one conversion depends on an argument type. Hand the literal + // text and the specs to the builder in cuda-device; its result is sized + // from the same data. + let part_lits: Vec = parts + .iter() + .map(|part| LitByteStr::new(part.as_bytes(), Span::call_site())) + .collect(); + let mut buf_len = parts.iter().map(String::len).sum::() + 1; + let mut spec_exprs: Vec = Vec::new(); + for conversion in &conversions { + match conversion { + Conversion::Fixed(spec) => { + buf_len += spec.len(); + let lit = LitByteStr::new(spec.as_bytes(), Span::call_site()); + spec_exprs.push(quote! { cuda_device::debug::__GpuPrintfSpec::Fixed(#lit) }); + } + Conversion::Typed { prefix, ty } => { + buf_len += prefix.len() + TYPED_SPEC_OVERHEAD; + let lit = LitByteStr::new(prefix.as_bytes(), Span::call_site()); + spec_exprs.push(quote! { + cuda_device::debug::__GpuPrintfSpec::Typed { + prefix: #lit, + format_char: <#ty as cuda_device::debug::GpuPrintfArg>::FORMAT_CHAR, + is_64bit: <#ty as cuda_device::debug::GpuPrintfArg>::IS_64BIT, + is_float: <#ty as cuda_device::debug::GpuPrintfArg>::IS_FLOAT, + } + }); + } + } + } + let fmt_ident = format_ident!("__fmt", span = Span::mixed_site()); - let expanded = quote! { + quote! { { - // Argument packing struct with C-compatible layout + // The format string of this call site, built once per + // instantiation. The sized array is the constant form the + // device importer accepts. + #[inline(always)] + fn __oxide_gpu_printf_format<#(#typed_params: cuda_device::debug::GpuPrintfArg),*>( + #(_: &#typed_params),* + ) -> &'static [u8; #buf_len] { + const { + &cuda_device::debug::__gpu_printf_format::<#buf_len>( + &[#(#part_lits),*], + &[#(#spec_exprs),*], + ) + } + } + + // Argument packing struct with C-compatible layout. Typed slots + // hold the `GpuPrintfArg` promoted type. #[repr(C)] - struct #struct_name { + struct __GpuPrintfArgs<#(#promoted_params),*> { #(#field_defs),* } - // Pack arguments - let __args = #struct_name { + #(#hoisted)* + + let #fmt_ident = __oxide_gpu_printf_format(#(&#typed_values),*); + let #args_ident = __GpuPrintfArgs { #(#field_inits),* }; - // Call vprintf intrinsic - // Use byte string literal directly (avoids str type which GPU doesn't support) cuda_device::debug::__gpu_vprintf( - #format_lit.as_ptr(), - &__args as *const #struct_name as *const u8 + #fmt_ident.as_ptr(), + &#args_ident as *const _ as *const u8 ) } - }; - - expanded + } } diff --git a/cuda-oxide/crates/cuda-macros/src/tests/launch.rs b/cuda-oxide/crates/cuda-macros/src/tests/launch.rs index eea6a710b3..9af3433200 100644 --- a/cuda-oxide/crates/cuda-macros/src/tests/launch.rs +++ b/cuda-oxide/crates/cuda-macros/src/tests/launch.rs @@ -16,7 +16,7 @@ use crate::launch_attrs::{ }; use quote::{format_ident, quote}; use reserved_oxide_symbols::INSTANTIATE_PREFIX; -use syn::{ItemFn, Stmt, parse_quote, visit_mut::VisitMut}; +use syn::{Expr, ItemFn, Stmt, parse_quote, spanned::Spanned, visit_mut::VisitMut}; #[test] fn generated_kernel_siblings_preserve_qualified_paths_and_generics() { @@ -307,6 +307,32 @@ fn unroll_n_on_for_loop_injects_factor_n_marker() { ); } +#[test] +fn unroll_marker_takes_the_position_of_its_attribute() { + // `parse_str` keeps real lines; `parse_quote!` gives every span line 0. + let cases = [ + ("fn k() {\n #[unroll]\n while true {}\n}", 2, 4), + ( + "fn k() {\n\n #[unroll(4)]\n for _ in 0..8 {}\n}", + 3, + 4, + ), + ("fn k() {\n #[unroll(2)]\n loop { break; }\n}", 2, 2), + ]; + for (source, line, column) in cases { + let mut func: ItemFn = syn::parse_str(source).expect("test source parses"); + rewrite_loop_unroll_attrs(&mut func).expect("valid unroll attribute"); + let marker = match &func.block.stmts[0] { + Stmt::Expr(Expr::While(l), _) => &l.body.stmts[0], + Stmt::Expr(Expr::ForLoop(l), _) => &l.body.stmts[0], + Stmt::Expr(Expr::Loop(l), _) => &l.body.stmts[0], + other => panic!("expected a loop, got {}", quote!(#other)), + }; + let start = marker.span().start(); + assert_eq!((start.line, start.column), (line, column), "{source}"); + } +} + #[test] fn unroll_on_bare_loop_injects_marker() { let func: ItemFn = parse_quote! { diff --git a/cuda-oxide/crates/cuda-macros/src/tests/mod.rs b/cuda-oxide/crates/cuda-macros/src/tests/mod.rs index 8e7851e3b5..4f378bd549 100644 --- a/cuda-oxide/crates/cuda-macros/src/tests/mod.rs +++ b/cuda-oxide/crates/cuda-macros/src/tests/mod.rs @@ -6,3 +6,5 @@ mod cuda_module; mod kernel; mod launch; +mod printf; +mod tile_export; diff --git a/cuda-oxide/crates/cuda-macros/src/tests/printf.rs b/cuda-oxide/crates/cuda-macros/src/tests/printf.rs new file mode 100644 index 0000000000..8cdcbb7a54 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/src/tests/printf.rs @@ -0,0 +1,290 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Expansion tests for `gpu_printf!`. +//! +//! Every assertion runs on the token text with spaces removed, so byte string +//! literals appear without their inner spaces too (`b" = "` is `b"="`). + +use crate::printf::{GpuPrintfInput, gpu_printf_impl}; +use syn::parse_quote; + +fn expanded(input: GpuPrintfInput) -> String { + gpu_printf_impl(input).to_string().replace(' ', "") +} + +/// Byte offset of `needle` in `haystack`, failing with the whole expansion. +fn position(haystack: &str, needle: &str) -> usize { + haystack + .find(needle) + .unwrap_or_else(|| panic!("expected `{needle}` in {haystack}")) +} + +fn assert_trait_route(expansion: &str) { + assert!( + expansion.contains("cuda_device::debug::GpuPrintfArg::promote(__val_"), + "{expansion}" + ); + assert!( + expansion.contains("cuda_device::debug::__gpu_printf_format::<"), + "{expansion}" + ); + assert!(expansion.contains(">::FORMAT_CHAR"), "{expansion}"); + // The macro's own casts wrap the argument in parentheses; a cast the + // user wrote (`value as u64`) renders without them. A runtime width is + // cast to `i32` on every route, so it is not in this list. + for cast in [")asi64", ")asu64", ")asf64", "%lld", "%llu", "%f"] { + assert!( + !expansion.contains(cast), + "default placeholder must not pick a conversion in the macro: `{cast}` in {expansion}" + ); + } +} + +#[test] +fn default_placeholders_take_the_trait_route_whatever_the_argument_looks_like() { + let inputs: [GpuPrintfInput; 8] = [ + parse_quote! { "{}", 42 }, + parse_quote! { "{}", 1.5 }, + parse_quote! { "{}", 18446744073709551615u64 }, + parse_quote! { "{}", 3.14159f32 }, + parse_quote! { "{}", true }, + parse_quote! { "{}", -x }, + parse_quote! { "{}", value as u64 }, + parse_quote! { "{}", data[tid] }, + ]; + for input in inputs { + let expansion = expanded(input); + assert_trait_route(&expansion); + assert!( + expansion.contains("__GpuPrintfSpec::Typed{prefix:b\"\","), + "{expansion}" + ); + } +} + +#[test] +fn width_zero_pad_flags_and_precision_ride_in_the_typed_prefix() { + let cases = [ + (parse_quote! { "{:8}", value }, "8"), + (parse_quote! { "{:08}", value }, "08"), + (parse_quote! { "{:.2}", value }, ".2"), + (parse_quote! { "{:-8.3}", value }, "-8.3"), + (parse_quote! { "{:+}", value }, "+"), + (parse_quote! { "{: }", value }, " "), + (parse_quote! { "{:#}", value }, "#"), + (parse_quote! { "{:*}", width, value }, "*"), + ]; + for (input, prefix) in cases { + let expansion = expanded(input); + assert_trait_route(&expansion); + let expected = format!( + "__GpuPrintfSpec::Typed{{prefix:b\"{}\",", + prefix.replace(' ', "") + ); + assert!(expansion.contains(&expected), "{expansion}"); + } +} + +#[test] +fn explicit_type_characters_keep_their_casts_and_the_literal_format() { + let cases = [ + ( + parse_quote! { "{:x}", 255u32 }, + "b\"%llx\\0\"", + "(255u32)asu64", + ), + ( + parse_quote! { "{:08X}", bits }, + "b\"%08llX\\0\"", + "(bits)asu64", + ), + ( + parse_quote! { "{:#o}", 64u32 }, + "b\"%#llo\\0\"", + "(64u32)asu64", + ), + (parse_quote! { "{:.2f}", pi }, "b\"%.2f\\0\"", "(pi)asf64"), + ( + parse_quote! { "{:e}", 1000.0f64 }, + "b\"%e\\0\"", + "(1000.0f64)asf64", + ), + ( + parse_quote! { "{:d}", value }, + "b\"%lld\\0\"", + "(value)asi64", + ), + ( + parse_quote! { "{:+d}", value }, + "b\"%+lld\\0\"", + "(value)asi64", + ), + (parse_quote! { "{:c}", 'a' }, "b\"%c\\0\"", "('a')asi32"), + (parse_quote! { "{:p}", ptr }, "b\"%llp\\0\"", "(ptr)asu64"), + ]; + for (input, format, cast) in cases { + let expansion = expanded(input); + assert!(expansion.contains(format), "{expansion}"); + assert!(expansion.contains(cast), "{expansion}"); + assert!( + !expansion.contains("__gpu_printf_format"), + "all-fixed call sites take the literal route: {expansion}" + ); + assert!(!expansion.contains("GpuPrintfArg::promote"), "{expansion}"); + } +} + +#[test] +fn mixed_call_sites_hand_fixed_and_typed_specs_to_the_builder_in_order() { + let expansion = expanded(parse_quote! { "{} = {:x} ({:.1})\n", count, bits, ratio }); + + assert!( + expansion.contains("&[b\"\",b\"=\",b\"(\",b\")\\n\"]"), + "{expansion}" + ); + let first = position( + &expansion, + "__GpuPrintfSpec::Typed{prefix:b\"\",format_char:<__T0as", + ); + let second = position(&expansion, "__GpuPrintfSpec::Fixed(b\"%llx\")"); + let third = position( + &expansion, + "__GpuPrintfSpec::Typed{prefix:b\".1\",format_char:<__T1as", + ); + assert!(first < second && second < third, "{expansion}"); + + assert!( + expansion.contains("__arg_0:cuda_device::debug::GpuPrintfArg::promote(__val_0)"), + "{expansion}" + ); + assert!(expansion.contains("__arg_1:__val_1"), "{expansion}"); + assert!( + expansion.contains("__arg_2:cuda_device::debug::GpuPrintfArg::promote(__val_2)"), + "{expansion}" + ); + assert!( + expansion.contains("__oxide_gpu_printf_format(&__val_0,&__val_2)"), + "{expansion}" + ); +} + +#[test] +fn the_format_buffer_is_sized_from_the_parts_and_specs() { + // "a" + "%" + prefix + "ll" + type character + "b" + NUL + let single = expanded(parse_quote! { "a{}b", value }); + assert!(single.contains("__gpu_printf_format::<7usize>"), "{single}"); + assert!(single.contains("->&'static[u8;7usize]"), "{single}"); + + // "100%% " is six bytes once the percent sign is doubled. + let percent = expanded(parse_quote! { "100% {}", value }); + assert!(percent.contains("b\"100%%\""), "{percent}"); + assert!( + percent.contains("__gpu_printf_format::<11usize>"), + "{percent}" + ); + + // "%llx" (4) + "%" + "08" + "ll" + type character (6) + NUL + let mixed = expanded(parse_quote! { "{:x}{:08}", bits, value }); + assert!(mixed.contains("__gpu_printf_format::<11usize>"), "{mixed}"); +} + +#[test] +fn arguments_are_evaluated_in_source_order_on_both_routes() { + let typed = expanded(parse_quote! { "{} [{:*}] {:x}\n", a(), w(), v(), h() }); + let lets = [ + "let__val_0=a();", + "let__val_1=(w())asi32;", + "let__val_2=v();", + "let__val_3=(h())asu64;", + "let__fmt=", + "let__args=", + ]; + let mut last = 0; + for needle in lets { + let at = position(&typed, needle); + assert!(at > last, "`{needle}` out of order in {typed}"); + last = at; + } + assert!(typed.contains("__width_1:__val_1"), "{typed}"); + assert!( + typed.contains("__arg_1:cuda_device::debug::GpuPrintfArg::promote(__val_2)"), + "{typed}" + ); + + let fixed = expanded(parse_quote! { "[{:*d}] {:x}\n", w(), v(), h() }); + let lets = [ + "let__val_0=(w())asi32;", + "let__val_1=(v())asi64;", + "let__val_2=(h())asu64;", + "let__args=", + ]; + let mut last = 0; + for needle in lets { + let at = position(&fixed, needle); + assert!(at > last, "`{needle}` out of order in {fixed}"); + last = at; + } + assert!(fixed.contains("b\"[%*lld]%llx\\n\\0\""), "{fixed}"); +} + +#[test] +fn literal_percent_signs_and_braces_are_escaped_on_both_routes() { + let typed = expanded(parse_quote! { "{{{}}} 100%\n", value }); + assert!(typed.contains("&[b\"{\",b\"}100%%\\n\"]"), "{typed}"); + + let fixed = expanded(parse_quote! { "{{{:d}}} 100%\n", value }); + assert!(fixed.contains("b\"{%lld}100%%\\n\\0\""), "{fixed}"); +} + +#[test] +fn a_call_without_placeholders_passes_a_null_argument_pointer() { + let expansion = expanded(parse_quote! { "hello 100%\n" }); + assert!(expansion.contains("b\"hello100%%\\n\\0\""), "{expansion}"); + assert!(expansion.contains("core::ptr::null()"), "{expansion}"); + assert!(!expansion.contains("__GpuPrintfArgs"), "{expansion}"); +} + +#[test] +fn string_placeholders_keep_the_literal_for_concat() { + let expansion = expanded(parse_quote! { "{:s}\n", "device" }); + assert!( + expansion + .contains("let__val_0={const__STR:&str=concat!(\"device\",\"\\0\");__STR.as_ptr()};"), + "{expansion}" + ); + assert!(expansion.contains("b\"%s\\n\\0\""), "{expansion}"); +} + +#[test] +fn runtime_precision_is_a_macro_error() { + let expansion = expanded(parse_quote! { "{:.*}", 2, pi }); + assert!(expansion.contains("compile_error"), "{expansion}"); + assert!( + expansion.contains("runtimeprecision`.*`isnotsupportedbyCUDAdeviceprintf"), + "{expansion}" + ); + + let with_type = expanded(parse_quote! { "{:.*f}", 2, pi }); + assert!(with_type.contains("compile_error"), "{with_type}"); +} + +#[test] +fn argument_count_mismatches_are_macro_errors() { + let missing = expanded(parse_quote! { "{} {}\n", 1 }); + assert!(missing.contains("compile_error"), "{missing}"); + assert!( + missing.contains("requires2argumentsbut1weresupplied"), + "{missing}" + ); + + // A runtime width takes an argument of its own. + let width = expanded(parse_quote! { "{:*}\n", 1 }); + assert!(width.contains("compile_error"), "{width}"); + assert!( + width.contains("requires2argumentsbut1weresupplied"), + "{width}" + ); +} diff --git a/cuda-oxide/crates/cuda-macros/src/tests/tile_export.rs b/cuda-oxide/crates/cuda-macros/src/tests/tile_export.rs new file mode 100644 index 0000000000..cbc39bb858 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/src/tests/tile_export.rs @@ -0,0 +1,385 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::tile_export::{TileExportArgs, expand_tile_export}; +use oxide_tile_export::{TileAccess, TileElem, TileExportDesc, TileSync}; +use proc_macro2::TokenStream as TokenStream2; +use quote::quote; +use reserved_oxide_symbols::tile_export_meta_symbol; +use syn::{ItemFn, parse_quote}; + +fn expand(args: TokenStream2, func: ItemFn) -> syn::Result { + let args: TileExportArgs = syn::parse2(args)?; + expand_tile_export(args, func).map(|tokens| tokens.to_string().replace(' ', "")) +} + +fn expand_err(args: TokenStream2, func: ItemFn) -> String { + match expand(args, func) { + Ok(expanded) => panic!("expected an error, got {expanded}"), + Err(err) => err.to_string(), + } +} + +/// Pull the embedded JSON back out of the expansion through the metadata +/// function's string literal. +fn embedded_desc(expanded: &str, symbol: &str) -> TileExportDesc { + let meta = tile_export_meta_symbol(symbol); + let start = expanded + .find(&format!("fn{meta}()")) + .unwrap_or_else(|| panic!("no metadata function `{meta}` in {expanded}")); + let rest = &expanded[start..]; + let open = rest.find('"').expect("string literal"); + let json: String = syn::parse_str::( + &rest[open..rest[open..].find("\"}").map(|end| open + end + 1).unwrap()], + ) + .expect("the metadata body is a string literal") + .value(); + TileExportDesc::from_json(&json).expect("embedded JSON parses") +} + +#[test] +fn scalar_in_and_returned_out_follow_the_table() { + let func: ItemFn = parse_quote! { + pub fn relu(x: f32) -> f32 { if x > 0.0 { x } else { 0.0 } } + }; + let expanded = expand(quote! { tile_shape = 128 }, func).unwrap(); + + // The original function survives unchanged. + assert!(expanded.contains("pubfnrelu(x:f32)->f32"), "{expanded}"); + // The wrapper is a `#[device]` extern "C" function on the native ABI. + assert!( + expanded.contains("#[::cuda_device::device]pubunsafeextern\"C\"fnrelu_impl("), + "{expanded}" + ); + assert!(expanded.contains(":f32,"), "{expanded}"); + assert!(expanded.contains(":*mutf32)"), "{expanded}"); + assert!(expanded.contains("=relu("), "{expanded}"); + + let desc = embedded_desc(&expanded, "relu_impl"); + desc.validate().unwrap(); + assert_eq!(desc.decl, "relu"); + assert_eq!(desc.symbol, "relu_impl"); + assert_eq!(desc.shape, vec![128]); + assert_eq!(desc.sync, TileSync::None); + assert!(desc.arch.is_empty()); + assert_eq!(desc.params.len(), 2); + assert_eq!(desc.params[0].name, "x"); + assert_eq!(desc.params[0].access, TileAccess::In); + assert!(desc.params[0].by_value); + assert_eq!(desc.params[1].name, "relu_out"); + assert_eq!(desc.params[1].access, TileAccess::Out); + assert_eq!(desc.params[1].elem, TileElem::F32); + assert!(!desc.params[1].by_value); +} + +#[test] +fn name_symbol_sync_and_arch_override_the_defaults() { + let func: ItemFn = parse_quote! { + pub fn scale(x: &mut f32, factor: f32) { *x *= factor; } + }; + let expanded = expand( + quote! { tile_shape = 128, name = "scale_in_place", symbol = "scale_impl_v2", sync = "warp", arch = ["sm_100", "sm_120"] }, + func, + ) + .unwrap(); + assert!( + expanded.contains("pubunsafeextern\"C\"fnscale_impl_v2("), + "{expanded}" + ); + assert!(expanded.contains(":*mutf32,"), "{expanded}"); + assert!(expanded.contains(":f32)"), "{expanded}"); + // `&mut T` is rebuilt from the pointer before the call. + assert!(expanded.contains("&mut*"), "{expanded}"); + + let desc = embedded_desc(&expanded, "scale_impl_v2"); + desc.validate().unwrap(); + assert_eq!(desc.decl, "scale_in_place"); + assert_eq!(desc.symbol, "scale_impl_v2"); + assert_eq!(desc.sync, TileSync::Warp); + assert_eq!(desc.arch, vec!["sm_100".to_string(), "sm_120".to_string()]); + assert_eq!(desc.params[0].access, TileAccess::InOut); + assert!(!desc.params[0].by_value); + assert_eq!(desc.params[1].access, TileAccess::In); + assert!(desc.params[1].by_value); +} + +#[test] +fn symbol_defaults_to_the_declaration_name() { + let func: ItemFn = parse_quote! { + pub fn scale(x: &mut f32) { *x *= 2.0; } + }; + let expanded = expand(quote! { tile_shape = 64, name = "double" }, func).unwrap(); + let desc = embedded_desc(&expanded, "double_impl"); + assert_eq!(desc.decl, "double"); + assert_eq!(desc.symbol, "double_impl"); +} + +#[test] +fn array_parameters_carry_their_footprint() { + let func: ItemFn = parse_quote! { + pub fn window(acc: &mut [i32; 4], bias: &[i32; 4]) { + let _ = (acc, bias); + } + }; + let expanded = expand(quote! { tile_shape = 64 }, func).unwrap(); + assert!(expanded.contains(":*muti32,"), "{expanded}"); + assert!(expanded.contains(":*consti32)"), "{expanded}"); + // Both arrays are rebuilt from their pointers before the call. + assert!(expanded.contains("as*mut[i32;4]"), "{expanded}"); + assert!(expanded.contains("as*const[i32;4]"), "{expanded}"); + + let desc = embedded_desc(&expanded, "window_impl"); + desc.validate().unwrap(); + let shapes: Vec<_> = desc + .params + .iter() + .map(|p| (p.name.as_str(), p.elem, p.access, p.footprint, p.by_value)) + .collect(); + assert_eq!( + shapes, + vec![ + ("acc", TileElem::I32, TileAccess::InOut, 4, false), + ("bias", TileElem::I32, TileAccess::In, 4, false), + ] + ); +} + +#[test] +fn scalar_references_and_halves_carry_footprint_one() { + let func: ItemFn = parse_quote! { + pub fn mix(scale: &'_ u16, half: f16, out: &mut u16) { + let _ = (scale, half, out); + } + }; + let expanded = expand(quote! { tile_shape = 64 }, func).unwrap(); + assert!(expanded.contains(":*constu16,"), "{expanded}"); + assert!(expanded.contains(":f16,"), "{expanded}"); + assert!(expanded.contains(":*mutu16)"), "{expanded}"); + + let desc = embedded_desc(&expanded, "mix_impl"); + desc.validate().unwrap(); + let shapes: Vec<_> = desc + .params + .iter() + .map(|p| (p.name.as_str(), p.elem, p.access, p.footprint, p.by_value)) + .collect(); + assert_eq!( + shapes, + vec![ + ("scale", TileElem::U16, TileAccess::In, 1, false), + ("half", TileElem::F16, TileAccess::In, 1, true), + ("out", TileElem::U16, TileAccess::InOut, 1, false), + ] + ); +} + +#[test] +fn mixed_footprints_are_rejected_at_the_differing_parameter() { + // The window needs 16 invocations of a 64-element tile, the scalar 64: + // one implementation cannot have both. + let func: ItemFn = parse_quote! { + pub fn window(acc: &mut [i32; 4], bias: &i32) { + let _ = (acc, bias); + } + }; + let err = expand_err(quote! { tile_shape = 64 }, func); + assert!( + err.contains("`acc` (footprint 4) and `bias` (footprint 1) differ"), + "{err}" + ); + // A by-value scalar is footprint 1 too, so it cannot sit beside a window. + let func: ItemFn = parse_quote! { + pub fn scale4(x: &mut [f32; 4], factor: f32) { + let _ = (x, factor); + } + }; + let err = expand_err(quote! { tile_shape = 128 }, func); + assert!( + err.contains("`x` (footprint 4) and `factor` (footprint 1) differ"), + "{err}" + ); +} + +#[test] +fn warp_sync_with_a_footprint_needs_whole_warps_of_footprints() { + let func: ItemFn = parse_quote! { + pub fn scale4(x: &mut [f32; 4]) { + let _ = x; + } + }; + // 64 is a multiple of 32, but a warp of 4-element footprints owns 128. + let err = expand_err(quote! { tile_shape = 64, sync = "warp" }, func.clone()); + assert!( + err.contains("footprint 4") && err.contains("multiple of 128") && err.contains("64"), + "{err}" + ); + expand(quote! { tile_shape = 128, sync = "warp" }, func.clone()).unwrap(); + // Without synchronization the footprint only has to divide the shape. + expand(quote! { tile_shape = 64 }, func).unwrap(); +} + +#[test] +fn primitive_module_paths_are_element_types() { + let func: ItemFn = parse_quote! { + pub fn add(a: core::primitive::f64, b: ::std::primitive::f64) -> f64 { a + b } + }; + let desc = embedded_desc( + &expand(quote! { tile_shape = 32 }, func).unwrap(), + "add_impl", + ); + assert!(desc.params.iter().all(|p| p.elem == TileElem::F64)); +} + +#[test] +fn the_out_parameter_avoids_a_user_parameter_of_the_same_name() { + let func: ItemFn = parse_quote! { + pub fn f(f_out: f32) -> f32 { f_out } + }; + let desc = embedded_desc(&expand(quote! { tile_shape = 32 }, func).unwrap(), "f_impl"); + assert_eq!(desc.params[1].name, "f_out_"); + desc.validate().unwrap(); +} + +#[test] +fn generics_self_and_async_are_rejected() { + let generic: ItemFn = parse_quote! { pub fn relu(x: T) -> T { x } }; + assert!(expand_err(quote! { tile_shape = 128 }, generic).contains("cannot be generic")); + + let lifetime: ItemFn = parse_quote! { pub fn relu<'a>(x: &'a f32) -> f32 { *x } }; + assert!(expand_err(quote! { tile_shape = 128 }, lifetime).contains("cannot be generic")); + + let method: ItemFn = parse_quote! { pub fn relu(self, x: f32) -> f32 { x } }; + assert!(expand_err(quote! { tile_shape = 128 }, method).contains("cannot take self")); + + let asynchronous: ItemFn = parse_quote! { pub async fn relu(x: f32) -> f32 { x } }; + assert!(expand_err(quote! { tile_shape = 128 }, asynchronous).contains("cannot be async")); +} + +#[test] +fn shape_is_required_positive_and_literal() { + let func: ItemFn = parse_quote! { pub fn relu(x: f32) -> f32 { x } }; + assert!(expand_err(quote! {}, func.clone()).contains("requires `tile_shape = N`")); + assert!(expand_err(quote! { tile_shape = 0 }, func.clone()).contains("positive integer")); + let err = syn::parse2::(quote! { tile_shape = N }) + .err() + .expect("a non-literal shape is a parse error") + .to_string(); + assert!(err.contains("integer literal"), "{err}"); + let err = syn::parse2::(quote! { tile_shape = 128, tile_shape = 64 }) + .err() + .expect("duplicate keys are rejected") + .to_string(); + assert!(err.contains("duplicate"), "{err}"); +} + +#[test] +fn unknown_and_malformed_arguments_are_rejected() { + let err = syn::parse2::(quote! { tile_shape = 128, foo = 1 }) + .err() + .expect("unknown keys are rejected") + .to_string(); + assert!( + err.contains("unknown #[tile_export] argument `foo`") + && err.contains("shape, name, symbol, sync, arch"), + "{err}" + ); + let err = syn::parse2::(quote! { tile_shape = 128, arch = "sm_120" }) + .err() + .expect("arch must be a list") + .to_string(); + assert!(err.contains("list of string literals"), "{err}"); + let err = syn::parse2::(quote! { tile_shape = 128, name = relu }) + .err() + .expect("name must be a string literal") + .to_string(); + assert!(err.contains("string literal"), "{err}"); + + let func: ItemFn = parse_quote! { pub fn relu(x: f32) -> f32 { x } }; + assert!( + expand_err(quote! { tile_shape = 128, sync = "cta" }, func.clone()) + .contains("must be \"none\" or \"warp\"") + ); + assert!( + expand_err(quote! { tile_shape = 128, symbol = "re-lu" }, func.clone()) + .contains("must be an identifier") + ); + assert!( + expand_err( + quote! { tile_shape = 128, name = "cuda_oxide_relu" }, + func.clone() + ) + .contains("reserved cuda-oxide prefix") + ); + assert!(expand_err(quote! { tile_shape = 128, arch = [""] }, func).contains("non-empty")); +} + +#[test] +fn unsupported_parameter_and_return_types_say_what_is_accepted() { + for func in [ + parse_quote! { pub fn bad(x: &[f32]) -> f32 { x[0] } }, + parse_quote! { pub fn bad(x: *mut f32) -> f32 { 0.0 } }, + parse_quote! { pub fn bad(x: bool) -> f32 { 0.0 } }, + parse_quote! { pub fn bad(x: usize) -> f32 { 0.0 } }, + parse_quote! { pub fn bad(x: &mut [f32]) -> f32 { 0.0 } }, + parse_quote! { pub fn bad(x: &[bool; 4]) -> f32 { 0.0 } }, + parse_quote! { pub fn bad(x: (f32, f32)) -> f32 { x.0 } }, + ] { + let err = expand_err(quote! { tile_shape = 128 }, func); + assert!( + err.contains( + "parameters must be `x: T`, `x: &T`, `x: &mut T`, `x: &[T; N]`, or \ + `x: &mut [T; N]`" + ), + "{err}" + ); + assert!( + err.contains("f16, f32, f64, i8, i16, i32, i64, u8, u16, u32, u64"), + "{err}" + ); + } + + let pattern: ItemFn = parse_quote! { pub fn bad((a, b): (f32, f32)) -> f32 { a } }; + assert!(expand_err(quote! { tile_shape = 128 }, pattern).contains("plain identifiers")); + + let const_len: ItemFn = parse_quote! { pub fn bad(x: &mut [f32; N]) {} }; + assert!(expand_err(quote! { tile_shape = 128 }, const_len).contains("integer literal length")); + + for func in [ + parse_quote! { pub fn bad(x: f32) -> &'static f32 { &0.0 } }, + parse_quote! { pub fn bad(x: f32) -> (f32, f32) { (x, x) } }, + parse_quote! { pub fn bad(x: f32) -> bool { x > 0.0 } }, + ] { + let err = expand_err(quote! { tile_shape = 128 }, func); + assert!(err.contains("must return a scalar"), "{err}"); + } +} + +#[test] +fn an_export_needs_an_output() { + let sink: ItemFn = parse_quote! { pub fn sink(x: f32, y: &f32) { let _ = (x, y); } }; + let err = expand_err(quote! { tile_shape = 128 }, sink); + assert!(err.contains("has no output"), "{err}"); + + // `()` is fine once something is written in place. + let in_place: ItemFn = parse_quote! { pub fn scale(x: &mut f32) -> () { *x *= 2.0; } }; + expand(quote! { tile_shape = 128 }, in_place).unwrap(); +} + +#[test] +fn warp_sync_needs_a_multiple_of_32() { + let func: ItemFn = parse_quote! { pub fn s(x: f32) -> f32 { x } }; + let err = expand_err(quote! { tile_shape = 100, sync = "warp" }, func.clone()); + assert!( + err.contains("multiple of 32") && err.contains("100"), + "{err}" + ); + expand(quote! { tile_shape = 96, sync = "warp" }, func).unwrap(); +} + +#[test] +fn reserved_function_names_are_rejected() { + let func: ItemFn = parse_quote! { pub fn cuda_oxide_relu(x: f32) -> f32 { x } }; + assert!(expand_err(quote! { tile_shape = 128 }, func).contains("reserved cuda-oxide")); +} diff --git a/cuda-oxide/crates/cuda-macros/src/tile_export.rs b/cuda-oxide/crates/cuda-macros/src/tile_export.rs new file mode 100644 index 0000000000..491e7d87b0 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/src/tile_export.rs @@ -0,0 +1,679 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! `#[tile_export]`: export a device function to cuTile through a tile +//! library. +//! +//! The attribute leaves the annotated function alone and adds two items next +//! to it: a `#[device] extern "C"` wrapper on the native SIMT ABI, and a +//! host-only `#[no_mangle]` metadata function whose body is the JSON +//! [`TileExportDesc`] of the export. The ABI rules live in `oxide-tile-export`; +//! this module maps Rust signatures onto them and emits the wrapper body. + +use crate::common::{internal_ident, reserved_name_error, track_codegen_environment}; +use oxide_tile_export::{ + CURRENT_VERSION, TileAccess, TileElem, TileExportDesc, TileExportError, TileParamDesc, + TileSync, is_identifier, +}; +use proc_macro::TokenStream; +use proc_macro2::{Span, TokenStream as TokenStream2}; +use quote::{ToTokens, quote}; +use reserved_oxide_symbols::{RESERVED_ROOT, tile_export_meta_symbol}; +use syn::{ + Expr, ExprArray, ExprLit, FnArg, Ident, ItemFn, Lit, LitInt, LitStr, Pat, ReturnType, + Signature, Token, Type, + parse::{Parse, ParseStream}, + parse_macro_input, +}; + +/// The Rust element types `#[tile_export]` accepts, as the diagnostics list +/// them. +const ELEMENT_TYPES: &str = "f16, f32, f64, i8, i16, i32, i64, u8, u16, u32, u64"; + +/// The attribute's parsed arguments. Every field is optional at parse time; +/// `expand_tile_export` reports the missing `tile_shape`. +#[derive(Default)] +pub(crate) struct TileExportArgs { + pub(crate) tile_shape: Option, + pub(crate) name: Option, + pub(crate) symbol: Option, + pub(crate) sync: Option, + pub(crate) arch: Option>, +} + +fn set_once(slot: &mut Option, key: &Ident, value: T) -> syn::Result<()> { + if slot.is_some() { + return Err(syn::Error::new( + key.span(), + format!("duplicate #[tile_export] argument `{key}`"), + )); + } + *slot = Some(value); + Ok(()) +} + +fn string_argument(key: &Ident, value: Expr) -> syn::Result { + match value { + Expr::Lit(ExprLit { + lit: Lit::Str(lit), .. + }) => Ok(lit), + other => Err(syn::Error::new_spanned( + other, + format!("#[tile_export] `{key}` takes a string literal such as `{key} = \"...\"`"), + )), + } +} + +impl Parse for TileExportArgs { + fn parse(input: ParseStream) -> syn::Result { + let mut args = Self::default(); + while !input.is_empty() { + let key: Ident = input.parse()?; + input.parse::()?; + match key.to_string().as_str() { + "tile_shape" => { + let value: Expr = input.parse()?; + let lit = match value { + Expr::Lit(ExprLit { + lit: Lit::Int(lit), .. + }) => lit, + other => { + return Err(syn::Error::new_spanned( + other, + "#[tile_export] `tile_shape` must be an integer literal such as \ + `tile_shape = 128`; a static tile library cannot ship a \ + parametric implementation, so a constant or expression is \ + not accepted", + )); + } + }; + set_once(&mut args.tile_shape, &key, lit)?; + } + "name" => { + let value = string_argument(&key, input.parse()?)?; + set_once(&mut args.name, &key, value)?; + } + "symbol" => { + let value = string_argument(&key, input.parse()?)?; + set_once(&mut args.symbol, &key, value)?; + } + "sync" => { + let value = string_argument(&key, input.parse()?)?; + set_once(&mut args.sync, &key, value)?; + } + "arch" => { + let array: ExprArray = input.parse().map_err(|err| { + syn::Error::new( + err.span(), + "#[tile_export] `arch` takes a list of string literals such as \ + `arch = [\"sm_100\", \"sm_120\"]`", + ) + })?; + let mut entries = Vec::with_capacity(array.elems.len()); + for elem in array.elems { + match elem { + Expr::Lit(ExprLit { + lit: Lit::Str(lit), .. + }) => entries.push(lit), + other => { + return Err(syn::Error::new_spanned( + other, + "#[tile_export] `arch` entries must be string literals \ + such as \"sm_120\"", + )); + } + } + } + set_once(&mut args.arch, &key, entries)?; + } + other => { + return Err(syn::Error::new( + key.span(), + format!( + "unknown #[tile_export] argument `{other}`; accepted arguments \ + are tile_shape, name, symbol, sync, arch" + ), + )); + } + } + if input.is_empty() { + break; + } + input.parse::()?; + } + Ok(args) + } +} + +pub(crate) fn tile_export_entry(attr: TokenStream, item: TokenStream) -> TokenStream { + track_codegen_environment(); + let args = parse_macro_input!(attr as TileExportArgs); + let input = parse_macro_input!(item as ItemFn); + match expand_tile_export(args, input) { + Ok(tokens) => tokens.into(), + Err(err) => err.to_compile_error().into(), + } +} + +/// How one Rust parameter travels across the native ABI. +enum ParamShape { + /// `x: T`: `in`, by value. + Scalar(TileElem), + /// `x: &T`: `in`, behind `*const T`. + Ref(TileElem), + /// `x: &[T; F]`: `in` with footprint `F`, behind `*const T`. + RefArray(TileElem, u32, LitInt), + /// `x: &mut T`: `inout`, behind `*mut T`. + RefMut(TileElem), + /// `x: &mut [T; F]`: `inout` with footprint `F`, behind `*mut T`. + RefMutArray(TileElem, u32, LitInt), +} + +impl ParamShape { + fn elem(&self) -> TileElem { + match self { + ParamShape::Scalar(elem) + | ParamShape::Ref(elem) + | ParamShape::RefArray(elem, _, _) + | ParamShape::RefMut(elem) + | ParamShape::RefMutArray(elem, _, _) => *elem, + } + } + + fn access(&self) -> TileAccess { + match self { + ParamShape::Scalar(_) | ParamShape::Ref(_) | ParamShape::RefArray(..) => TileAccess::In, + ParamShape::RefMut(_) | ParamShape::RefMutArray(..) => TileAccess::InOut, + } + } + + fn footprint(&self) -> u32 { + match self { + ParamShape::RefArray(_, footprint, _) | ParamShape::RefMutArray(_, footprint, _) => { + *footprint + } + _ => 1, + } + } + + fn by_value(&self) -> bool { + matches!(self, ParamShape::Scalar(_)) + } +} + +/// A type as the user would write it, for diagnostics: token printing puts +/// spaces around `&`, `;`, and `::` that no one types. +fn type_text(ty: &Type) -> String { + ty.to_token_stream() + .to_string() + .replace("& mut ", "&mut ") + .replace("& ", "&") + .replace(" ;", ";") + .replace(" ::", "::") + .replace(":: ", "::") + .replace(" <", "<") + .replace("< ", "<") + .replace(" >", ">") + .replace(" ,", ",") +} + +/// Peel `(T)` and macro-generated groups so a type written through another +/// macro classifies like one written by hand. +fn unwrap_type(mut ty: &Type) -> &Type { + loop { + match ty { + Type::Paren(paren) => ty = &paren.elem, + Type::Group(group) => ty = &group.elem, + _ => return ty, + } + } +} + +/// The element type a scalar type path names, if it is one of the eleven. +/// +/// Accepts the bare name (`f32`) and the primitive module paths +/// (`core::primitive::f32`, `::std::primitive::f32`). Anything with generic +/// arguments or any other path is not an element type. +fn scalar_elem(ty: &Type) -> Option { + let Type::Path(path) = unwrap_type(ty) else { + return None; + }; + if path.qself.is_some() { + return None; + } + let segments = &path.path.segments; + if segments.iter().any(|segment| !segment.arguments.is_none()) { + return None; + } + let last = segments.last()?; + match segments.len() { + 1 => {} + 3 => { + let crate_name = segments[0].ident.to_string(); + if !matches!(crate_name.as_str(), "core" | "std") || segments[1].ident != "primitive" { + return None; + } + } + _ => return None, + } + TileElem::from_rust_name(&last.ident.to_string()) +} + +fn classify_param_type(ty: &Type) -> syn::Result { + let unsupported = || { + syn::Error::new_spanned( + ty, + format!( + "#[tile_export] parameters must be `x: T`, `x: &T`, `x: &mut T`, \ + `x: &[T; N]`, or `x: &mut [T; N]` with T one of {ELEMENT_TYPES} and N an \ + integer literal; found `{}`. Slices, raw pointers, tuples, structs, \ + generics, and other element types are not accepted", + type_text(ty) + ), + ) + }; + if let Some(elem) = scalar_elem(ty) { + return Ok(ParamShape::Scalar(elem)); + } + let Type::Reference(reference) = unwrap_type(ty) else { + return Err(unsupported()); + }; + let mutable = reference.mutability.is_some(); + let pointee = unwrap_type(&reference.elem); + if let Some(elem) = scalar_elem(pointee) { + return Ok(if mutable { + ParamShape::RefMut(elem) + } else { + ParamShape::Ref(elem) + }); + } + let Type::Array(array) = pointee else { + return Err(unsupported()); + }; + let Some(elem) = scalar_elem(&array.elem) else { + return Err(unsupported()); + }; + let Expr::Lit(ExprLit { + lit: Lit::Int(len), .. + }) = &array.len + else { + return Err(syn::Error::new_spanned( + &array.len, + "#[tile_export] array parameters need an integer literal length such as \ + `&[f32; 4]` or `&mut [f32; 4]`; the footprint must be known to the macro, so \ + a constant or expression is not accepted", + )); + }; + let footprint: u32 = len.base10_parse()?; + if footprint == 0 { + return Err(syn::Error::new_spanned( + len, + "#[tile_export] array parameters need at least one element", + )); + } + Ok(if mutable { + ParamShape::RefMutArray(elem, footprint, len.clone()) + } else { + ParamShape::RefArray(elem, footprint, len.clone()) + }) +} + +fn classify_return_type(output: &ReturnType) -> syn::Result> { + let ty = match output { + ReturnType::Default => return Ok(None), + ReturnType::Type(_, ty) => ty, + }; + if let Type::Tuple(tuple) = unwrap_type(ty) + && tuple.elems.is_empty() + { + return Ok(None); + } + scalar_elem(ty).map(Some).ok_or_else(|| { + syn::Error::new_spanned( + ty, + format!( + "#[tile_export] functions must return a scalar (one of {ELEMENT_TYPES}) or \ + nothing; found `{}`. References, pointers, arrays, tuples, and other \ + element types are not accepted as return types", + type_text(ty) + ), + ) + }) +} + +/// Reject a declaration or symbol name that cannot be a Tile IR symbol or +/// that lands in the reserved namespace. +fn check_symbol_literal(lit: &LitStr, what: &str) -> syn::Result { + let value = lit.value(); + if !is_identifier(&value) { + return Err(syn::Error::new_spanned( + lit, + format!( + "#[tile_export] {what} `{value}` must be an identifier (ASCII letters, digits \ + and `_`, not starting with a digit)" + ), + )); + } + if value.starts_with(RESERVED_ROOT) { + return Err(syn::Error::new_spanned( + lit, + format!( + "#[tile_export] {what} `{value}` starts with the reserved cuda-oxide prefix \ + `{RESERVED_ROOT}`; choose another name" + ), + )); + } + Ok(value) +} + +/// A shared-rule violation as a diagnostic, pointed at the part of the +/// signature that breaks the rule: the parameter whose footprint differs, the +/// shape a warp export cannot split, or the function name for the rest. +fn validation_error( + err: TileExportError, + sig: &Signature, + shape_lit: &LitInt, + fn_ident: &Ident, +) -> syn::Error { + let message = format!("#[tile_export]: {err}"); + match &err { + TileExportError::MixedFootprints { other, .. } => { + let offending = sig.inputs.iter().find_map(|arg| match arg { + FnArg::Typed(pat_type) => match &*pat_type.pat { + Pat::Ident(pat_ident) if pat_ident.ident == other => Some(&pat_type.ty), + _ => None, + }, + FnArg::Receiver(_) => None, + }); + match offending { + Some(ty) => syn::Error::new_spanned(ty, message), + None => syn::Error::new(fn_ident.span(), message), + } + } + TileExportError::WarpFootprintShape { .. } | TileExportError::WarpShape { .. } => { + syn::Error::new_spanned(shape_lit, message) + } + _ => syn::Error::new(fn_ident.span(), message), + } +} + +/// Expand `#[tile_export(args)] fn ...` into the original function, the +/// `#[device] extern "C"` wrapper, and the metadata function. +pub(crate) fn expand_tile_export(args: TileExportArgs, input: ItemFn) -> syn::Result { + let sig = &input.sig; + let fn_ident = sig.ident.clone(); + if let Some(err) = reserved_name_error(&fn_ident) { + return Err(err); + } + if !sig.generics.params.is_empty() || sig.generics.where_clause.is_some() { + return Err(syn::Error::new_spanned( + &sig.generics, + "#[tile_export] functions cannot be generic: lifetime, type, and const \ + parameters are not accepted because a static tile library needs one concrete \ + symbol per export", + )); + } + if let Some(asyncness) = &sig.asyncness { + return Err(syn::Error::new_spanned( + asyncness, + "#[tile_export] functions cannot be async; only plain `fn` items run on the device", + )); + } + if let Some(variadic) = &sig.variadic { + return Err(syn::Error::new_spanned( + variadic, + "#[tile_export] functions cannot be variadic", + )); + } + + let Some(shape_lit) = args.tile_shape else { + return Err(syn::Error::new( + fn_ident.span(), + "#[tile_export] requires `tile_shape = N`, the 1-D tile extent; a static tile library \ + cannot ship a parametric implementation, so every export names one extent", + )); + }; + let shape: i64 = shape_lit.base10_parse()?; + if shape <= 0 { + return Err(syn::Error::new_spanned( + &shape_lit, + format!("#[tile_export] `tile_shape` must be a positive integer, found {shape}"), + )); + } + let sync = match &args.sync { + None => TileSync::None, + Some(lit) => match lit.value().as_str() { + "none" => TileSync::None, + "warp" => TileSync::Warp, + other => { + return Err(syn::Error::new_spanned( + lit, + format!( + "#[tile_export] `sync` must be \"none\" or \"warp\", found \"{other}\"" + ), + )); + } + }, + }; + // The footprint-1 form of the warp rule, reported before the parameters + // are read so the message points at the shape. The footprint-aware form + // comes from `validate` once the parameters are known. + if sync == TileSync::Warp && shape % 32 != 0 { + return Err(syn::Error::new_spanned( + &shape_lit, + format!( + "#[tile_export] sync = \"warp\" needs a tile_shape that is a multiple of 32 so \ + every invocation group fills whole warps, found {shape}" + ), + )); + } + let decl = match &args.name { + Some(lit) => check_symbol_literal(lit, "name")?, + None => fn_ident.to_string(), + }; + let symbol = match &args.symbol { + Some(lit) => check_symbol_literal(lit, "symbol")?, + None => format!("{decl}_impl"), + }; + let arch: Vec = args + .arch + .iter() + .flatten() + .map(|lit| { + let value = lit.value(); + if value.trim().is_empty() { + Err(syn::Error::new_spanned( + lit, + "#[tile_export] `arch` entries must be non-empty names such as \"sm_120\"", + )) + } else { + Ok(value) + } + }) + .collect::>()?; + + // Parameters: plain identifiers with one of the four accepted types. + let mut shapes = Vec::with_capacity(sig.inputs.len()); + let mut params = Vec::with_capacity(sig.inputs.len() + 1); + let mut types = Vec::with_capacity(sig.inputs.len()); + for arg in &sig.inputs { + let FnArg::Typed(pat_type) = arg else { + return Err(syn::Error::new_spanned( + arg, + "#[tile_export] functions cannot take self; export a free function", + )); + }; + if let Some(attr) = pat_type.attrs.first() { + return Err(syn::Error::new_spanned( + attr, + "#[tile_export] parameters cannot carry attributes", + )); + } + let Pat::Ident(pat_ident) = &*pat_type.pat else { + return Err(syn::Error::new_spanned( + &pat_type.pat, + "#[tile_export] parameters must be plain identifiers such as `x: f32`; \ + patterns are not accepted because each parameter names one Tile IR argument", + )); + }; + if pat_ident.subpat.is_some() || pat_ident.by_ref.is_some() { + return Err(syn::Error::new_spanned( + pat_ident, + "#[tile_export] parameters must be plain identifiers such as `x: f32`; \ + `ref` bindings and subpatterns are not accepted", + )); + } + let name = pat_ident.ident.to_string(); + if !is_identifier(&name) { + return Err(syn::Error::new_spanned( + &pat_ident.ident, + format!( + "#[tile_export] parameter `{name}` must be an ASCII identifier without \ + the `r#` prefix; it names a Tile IR argument" + ), + )); + } + let shape = classify_param_type(&pat_type.ty)?; + params.push(TileParamDesc { + name, + elem: shape.elem(), + access: shape.access(), + footprint: shape.footprint(), + by_value: shape.by_value(), + }); + types.push((*pat_type.ty).clone()); + shapes.push(shape); + } + + let return_elem = classify_return_type(&sig.output)?; + if return_elem.is_none() && !params.iter().any(|param| param.access.is_output()) { + return Err(syn::Error::new_spanned( + sig, + "#[tile_export] function has no output: return a scalar value or take at least \ + one `&mut T` or `&mut [T; N]` parameter, because a tile library declaration's \ + results are its `out` and `inout` arguments", + )); + } + if let Some(elem) = return_elem { + // The declaration argument that carries the return value. Keep the + // conventional `_out` unless a parameter already took it. + let mut out_name = format!("{fn_ident}_out"); + while params.iter().any(|param| param.name == out_name) { + out_name.push('_'); + } + params.push(TileParamDesc { + name: out_name, + elem, + access: TileAccess::Out, + footprint: 1, + by_value: false, + }); + } + + let desc = TileExportDesc { + version: CURRENT_VERSION, + decl, + symbol, + shape: vec![shape], + params, + sync, + arch, + }; + desc.validate() + .map_err(|err| validation_error(err, sig, &shape_lit, &fn_ident))?; + + // The wrapper: pointer parameters are turned back into the references + // the original function takes, the return value is stored through the + // trailing out pointer. + let mut abi_inputs = Vec::with_capacity(shapes.len() + 1); + let mut call_args = Vec::with_capacity(shapes.len()); + for (index, (shape, ty)) in shapes.iter().zip(&types).enumerate() { + let ident = internal_ident(&format!("__cuda_oxide_tile_arg_{index}")); + match shape { + ParamShape::Scalar(_) => { + abi_inputs.push(quote! { #ident: #ty }); + call_args.push(quote! { #ident }); + } + ParamShape::Ref(_) => { + let Type::Reference(reference) = unwrap_type(ty) else { + unreachable!("classified as a reference"); + }; + let pointee = &reference.elem; + abi_inputs.push(quote! { #ident: *const #pointee }); + call_args.push(quote! { unsafe { &*#ident } }); + } + ParamShape::RefMut(_) => { + let Type::Reference(reference) = unwrap_type(ty) else { + unreachable!("classified as a reference"); + }; + let pointee = &reference.elem; + abi_inputs.push(quote! { #ident: *mut #pointee }); + call_args.push(quote! { unsafe { &mut *#ident } }); + } + ParamShape::RefArray(_, _, len) => { + let Type::Reference(reference) = unwrap_type(ty) else { + unreachable!("classified as a reference"); + }; + let Type::Array(array) = unwrap_type(&reference.elem) else { + unreachable!("classified as an array reference"); + }; + let elem_ty = &array.elem; + abi_inputs.push(quote! { #ident: *const #elem_ty }); + call_args.push(quote! { unsafe { &*(#ident as *const [#elem_ty; #len]) } }); + } + ParamShape::RefMutArray(_, _, len) => { + let Type::Reference(reference) = unwrap_type(ty) else { + unreachable!("classified as a reference"); + }; + let Type::Array(array) = unwrap_type(&reference.elem) else { + unreachable!("classified as an array reference"); + }; + let elem_ty = &array.elem; + abi_inputs.push(quote! { #ident: *mut #elem_ty }); + call_args.push(quote! { unsafe { &mut *(#ident as *mut [#elem_ty; #len]) } }); + } + } + } + let call = quote! { #fn_ident(#(#call_args),*) }; + let body = match &sig.output { + ReturnType::Type(_, ret_ty) if return_elem.is_some() => { + let out_ident = internal_ident("__cuda_oxide_tile_out"); + let result_ident = internal_ident("__cuda_oxide_tile_result"); + abi_inputs.push(quote! { #out_ident: *mut #ret_ty }); + quote! { + let #result_ident: #ret_ty = #call; + unsafe { *#out_ident = #result_ident; } + } + } + _ => quote! { #call; }, + }; + + let wrapper_ident = Ident::new(&desc.symbol, fn_ident.span()); + let meta_ident = Ident::new(&tile_export_meta_symbol(&desc.symbol), fn_ident.span()); + let json = LitStr::new(&desc.to_json(), Span::call_site()); + let wrapper_doc = format!( + "Tile library implementation of `{}`, on the native SIMT ABI.", + desc.decl + ); + + Ok(quote! { + #input + + #[doc = #wrapper_doc] + #[doc(hidden)] + #[allow(non_snake_case, clippy::missing_safety_doc, improper_ctypes_definitions)] + #[::cuda_device::device] + pub unsafe extern "C" fn #wrapper_ident(#(#abi_inputs),*) { + #body + } + + #[doc(hidden)] + #[allow(non_snake_case)] + #[unsafe(no_mangle)] + pub fn #meta_ident() -> &'static str { + #json + } + }) +} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_requires_unresolved_constant.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_requires_unresolved_constant.stderr index 2d44609cbf..f713660dea 100644 --- a/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_requires_unresolved_constant.stderr +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/launch_contract_requires_unresolved_constant.stderr @@ -1,14 +1,12 @@ error[E0425]: cannot find value `TIEL` in this scope --> tests/compile_fail/launch_contract_requires_unresolved_constant.rs:16:40 | +10 | const TILE: usize = 36; + | ----------------------- similarly named constant `TILE` defined here +... 16 | requires = (input.len() >= n * TIEL), - | ^^^^ not found in this scope - | -note: similarly named constant `TILE` defined here - --> tests/compile_fail/launch_contract_requires_unresolved_constant.rs:10:5 + | ^^^^ | -10 | const TILE: usize = 36; - | ^^^^^^^^^^^^^^^^^^^^^^^ help: a constant with a similar name exists | 16 - requires = (input.len() >= n * TIEL), 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 a5f714062b..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 @@ -25,9 +25,7 @@ note: erroneous constant encountered | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ note: the above error was encountered while instantiating `fn __unroll_config::<1>` - --> tests/compile_fail/policy_config_invalid_unroll.rs:12:1 + --> tests/compile_fail/policy_config_invalid_unroll.rs:15:5 | -12 | #[kernel] - | ^^^^^^^^^ - | - = note: this note originates in the attribute macro `kernel` (in Nightly builds, run with -Z macro-backtrace for more info) +15 | #[unroll(Invalid::UNROLL)] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^ 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_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/tile_export_generic.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_generic.rs new file mode 100644 index 0000000000..aa9f403c35 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_generic.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +#[tile_export(tile_shape = 128)] +pub fn relu(x: T) -> T { + x +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_generic.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_generic.stderr new file mode 100644 index 0000000000..2b3433b290 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_generic.stderr @@ -0,0 +1,5 @@ +error: #[tile_export] functions cannot be generic: lifetime, type, and const parameters are not accepted because a static tile library needs one concrete symbol per export + --> tests/compile_fail/tile_export_generic.rs:7:12 + | +7 | pub fn relu(x: T) -> T { + | ^^^^^^^^^ diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_missing_shape.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_missing_shape.rs new file mode 100644 index 0000000000..d92422c42a --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_missing_shape.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +#[tile_export] +pub fn relu(x: f32) -> f32 { + x +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_missing_shape.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_missing_shape.stderr new file mode 100644 index 0000000000..c6bda38223 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_missing_shape.stderr @@ -0,0 +1,5 @@ +error: #[tile_export] requires `tile_shape = N`, the 1-D tile extent; a static tile library cannot ship a parametric implementation, so every export names one extent + --> tests/compile_fail/tile_export_missing_shape.rs:7:8 + | +7 | pub fn relu(x: f32) -> f32 { + | ^^^^ diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_mixed_footprints.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_mixed_footprints.rs new file mode 100644 index 0000000000..9bb6edc906 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_mixed_footprints.rs @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +#[tile_export(tile_shape = 64)] +pub fn window(acc: &mut [i32; 4], bias: &i32) { + for value in acc.iter_mut() { + *value += *bias; + } +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_mixed_footprints.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_mixed_footprints.stderr new file mode 100644 index 0000000000..323f78f903 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_mixed_footprints.stderr @@ -0,0 +1,5 @@ +error: #[tile_export]: parameters `acc` (footprint 4) and `bias` (footprint 1) differ; every parameter of one export shares the invocation count, so all footprints must be equal + --> tests/compile_fail/tile_export_mixed_footprints.rs:7:41 + | +7 | pub fn window(acc: &mut [i32; 4], bias: &i32) { + | ^^^^ diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_no_outputs.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_no_outputs.rs new file mode 100644 index 0000000000..a654191668 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_no_outputs.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +#[tile_export(tile_shape = 128)] +pub fn sink(x: f32, y: &f32) { + let _ = (x, y); +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_no_outputs.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_no_outputs.stderr new file mode 100644 index 0000000000..c58c08c60d --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_no_outputs.stderr @@ -0,0 +1,5 @@ +error: #[tile_export] function has no output: return a scalar value or take at least one `&mut T` or `&mut [T; N]` parameter, because a tile library declaration's results are its `out` and `inout` arguments + --> tests/compile_fail/tile_export_no_outputs.rs:7:5 + | +7 | pub fn sink(x: f32, y: &f32) { + | ^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_outputs_split.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_outputs_split.rs new file mode 100644 index 0000000000..e7220abc4b --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_outputs_split.rs @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +// [inout, in, out]: the appended return slot is separated from `a` by `b`. +#[tile_export(tile_shape = 128)] +pub fn split(a: &mut f32, b: f32) -> f32 { + *a *= b; + *a +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_outputs_split.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_outputs_split.stderr new file mode 100644 index 0000000000..a23ce382d6 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_outputs_split.stderr @@ -0,0 +1,5 @@ +error: #[tile_export]: `out` and `inout` parameters must be one contiguous run, but `b` is `in` and sits between `a` and `split_out`; move `in` parameters before the `&mut` ones + --> tests/compile_fail/tile_export_outputs_split.rs:8:8 + | +8 | pub fn split(a: &mut f32, b: f32) -> f32 { + | ^^^^^ diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unknown_argument.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unknown_argument.rs new file mode 100644 index 0000000000..4db9fd79ef --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unknown_argument.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +#[tile_export(tile_shape = 128, shap = 64)] +pub fn relu(x: f32) -> f32 { + x +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unknown_argument.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unknown_argument.stderr new file mode 100644 index 0000000000..cac0df871b --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unknown_argument.stderr @@ -0,0 +1,5 @@ +error: unknown #[tile_export] argument `shap`; accepted arguments are tile_shape, name, symbol, sync, arch + --> tests/compile_fail/tile_export_unknown_argument.rs:6:33 + | +6 | #[tile_export(tile_shape = 128, shap = 64)] + | ^^^^ diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unsupported_param.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unsupported_param.rs new file mode 100644 index 0000000000..9c7dd58627 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unsupported_param.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +#[tile_export(tile_shape = 128)] +pub fn first(x: &[f32]) -> f32 { + x[0] +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unsupported_param.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unsupported_param.stderr new file mode 100644 index 0000000000..bd7f56cfa5 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_unsupported_param.stderr @@ -0,0 +1,5 @@ +error: #[tile_export] parameters must be `x: T`, `x: &T`, `x: &mut T`, `x: &[T; N]`, or `x: &mut [T; N]` with T one of f16, f32, f64, i8, i16, i32, i64, u8, u16, u32, u64 and N an integer literal; found `&[f32]`. Slices, raw pointers, tuples, structs, generics, and other element types are not accepted + --> tests/compile_fail/tile_export_unsupported_param.rs:7:17 + | +7 | pub fn first(x: &[f32]) -> f32 { + | ^^^^^^ diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_footprint.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_footprint.rs new file mode 100644 index 0000000000..f52d1b2137 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_footprint.rs @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +#[tile_export(tile_shape = 64, sync = "warp")] +pub fn scale4(x: &mut [f32; 4]) { + for value in x.iter_mut() { + *value *= 2.0; + } +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_footprint.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_footprint.stderr new file mode 100644 index 0000000000..89fa47f5b3 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_footprint.stderr @@ -0,0 +1,5 @@ +error: #[tile_export]: sync = warp with footprint 4 needs a shape that is a multiple of 128 (32 lanes times the footprint), found 64 + --> tests/compile_fail/tile_export_warp_footprint.rs:6:28 + | +6 | #[tile_export(tile_shape = 64, sync = "warp")] + | ^^ diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_shape.rs b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_shape.rs new file mode 100644 index 0000000000..cf98bcafec --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_shape.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use cuda_macros::tile_export; + +#[tile_export(tile_shape = 100, sync = "warp")] +pub fn partial_warp(x: f32) -> f32 { + x +} + +fn main() {} diff --git a/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_shape.stderr b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_shape.stderr new file mode 100644 index 0000000000..0987548c1e --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/compile_fail/tile_export_warp_shape.stderr @@ -0,0 +1,5 @@ +error: #[tile_export] sync = "warp" needs a tile_shape that is a multiple of 32 so every invocation group fills whole warps, found 100 + --> tests/compile_fail/tile_export_warp_shape.rs:6:28 + | +6 | #[tile_export(tile_shape = 100, sync = "warp")] + | ^^^ 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 f7fdc1eff6..4a78f429b9 100644 --- a/cuda-oxide/crates/cuda-macros/tests/device-only/Cargo.lock +++ b/cuda-oxide/crates/cuda-macros/tests/device-only/Cargo.lock @@ -13,10 +13,11 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", - "syn", + "syn 2.0.119", ] [[package]] @@ -27,6 +28,27 @@ dependencies = [ "cuda-macros", ] +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -49,6 +71,49 @@ dependencies = [ name = "reserved-oxide-symbols" version = "0.2.1" +[[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.9", +] + +[[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 = "syn" version = "2.0.119" @@ -60,8 +125,45 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +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.9", +] + [[package]] name = "unicode-ident" version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[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/cuda-macros/tests/macro_guard.rs b/cuda-oxide/crates/cuda-macros/tests/macro_guard.rs index c25731140c..6ab0f6691c 100644 --- a/cuda-oxide/crates/cuda-macros/tests/macro_guard.rs +++ b/cuda-oxide/crates/cuda-macros/tests/macro_guard.rs @@ -9,6 +9,7 @@ fn macro_guards() { let t = trybuild::TestCases::new(); t.pass("tests/pass/const_generic_hygiene.rs"); + t.pass("tests/pass/unroll_literal_factor_lint_clean.rs"); t.pass("tests/pass/cuda_module_inline_namespaces.rs"); t.pass("tests/pass/cuda_module_generated_by_macro_rules.rs"); t.compile_fail("tests/compile_fail/kernel_reserved_name.rs"); 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..29cd8e0a86 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); // 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/tile_export_forms.rs b/cuda-oxide/crates/cuda-macros/tests/pass/tile_export_forms.rs new file mode 100644 index 0000000000..38a53fc6db --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/pass/tile_export_forms.rs @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Every parameter and return form `#[tile_export]` accepts, expanded +//! against the real `cuda_device::device` so the generated wrapper +//! type-checks. + +#![allow(dead_code)] +#![feature(f16)] + +use cuda_device::tile_export; + +#[tile_export(tile_shape = 128)] +pub fn relu(x: f32) -> f32 { + if x > 0.0 { x } else { 0.0 } +} + +#[tile_export(tile_shape = 128, name = "scale_in_place", symbol = "scale_in_place_impl")] +pub fn scale(x: &mut f32, factor: f32) { + *x *= factor; +} + +#[tile_export(tile_shape = 256, sync = "warp", arch = ["sm_100", "sm_120"])] +pub fn warp_sum(x: f32) -> f32 { + x +} + +// Every parameter of one export shares its footprint, so the bias is a +// 4-element window too. +#[tile_export(tile_shape = 64)] +pub fn window(acc: &mut [i32; 4], bias: &[i32; 4]) { + for (value, bias) in acc.iter_mut().zip(bias) { + *value += *bias; + } +} + +// A warp of 4-element footprints owns 128 elements; 256 holds two. +#[tile_export(tile_shape = 256, sync = "warp")] +pub fn warp_scale4(x: &mut [f32; 4], factor: &[f32; 4]) { + for (value, factor) in x.iter_mut().zip(factor) { + *value *= *factor; + } +} + +#[tile_export(tile_shape = 32)] +pub fn half_add(a: f16, b: f16) -> f16 { + a + b +} + +#[tile_export(tile_shape = 32)] +pub fn widen(x: u8, y: i64) -> u64 { + (x as u64).wrapping_add(y as u64) +} + +mod nested { + use cuda_device::tile_export; + + #[tile_export(tile_shape = 128, name = "nested_negate")] + pub fn negate(x: &mut f64) { + *x = -*x; + } +} + +fn main() { + // The original functions stay ordinary Rust functions. + assert_eq!(relu(-1.0), 0.0); + let mut value = 2.0; + scale(&mut value, 3.0); + assert_eq!(value, 6.0); +} diff --git a/cuda-oxide/crates/cuda-macros/tests/pass/unroll_literal_factor_lint_clean.rs b/cuda-oxide/crates/cuda-macros/tests/pass/unroll_literal_factor_lint_clean.rs new file mode 100644 index 0000000000..e44e921cc5 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/pass/unroll_literal_factor_lint_clean.rs @@ -0,0 +1,45 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// A literal `#[unroll(N)]` expands to a marker call whose const argument sits +// in braces. Those braces must keep the macro's own span. On the attribute's +// span rustc treats them as written by the user, and its `unused_braces` lint +// then reports every literal factor in the user's crate, with a fix-it that +// deletes the attribute. This crate denies that lint, so the build fails if +// the braces ever take the attribute's span again. +#![deny(unused_braces)] + +use cuda_device::{device, kernel}; + +#[device] +fn count_down(mut remaining: u32) -> u32 { + let mut steps = 0; + #[unroll] + loop { + if remaining == 0 { + break; + } + remaining -= 1; + steps += 1; + } + steps +} + +#[kernel] +fn literal_factors(limit: u32) { + let mut total = 0; + #[unroll(4)] + for i in 0..8 { + total += i; + } + let mut index = 0; + #[unroll(2)] + while index < limit { + index += 1; + } + let _ = (total, count_down(limit)); +} + +fn main() { + assert_eq!(count_down(3), 3); +} diff --git a/cuda-oxide/crates/cuda-macros/tests/tile_export_guard.rs b/cuda-oxide/crates/cuda-macros/tests/tile_export_guard.rs new file mode 100644 index 0000000000..c220ffd819 --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/tile_export_guard.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `#[tile_export]` signature and argument contracts: every accepted form +//! expands against the real `cuda_device::device`, and each rejected form +//! fails with a macro diagnostic that says what is accepted instead. + +#[test] +fn tile_export_contracts() { + let t = trybuild::TestCases::new(); + t.pass("tests/pass/tile_export_forms.rs"); + t.compile_fail("tests/compile_fail/tile_export_generic.rs"); + t.compile_fail("tests/compile_fail/tile_export_missing_shape.rs"); + t.compile_fail("tests/compile_fail/tile_export_unsupported_param.rs"); + t.compile_fail("tests/compile_fail/tile_export_no_outputs.rs"); + t.compile_fail("tests/compile_fail/tile_export_warp_shape.rs"); + t.compile_fail("tests/compile_fail/tile_export_warp_footprint.rs"); + t.compile_fail("tests/compile_fail/tile_export_mixed_footprints.rs"); + t.compile_fail("tests/compile_fail/tile_export_unknown_argument.rs"); + t.compile_fail("tests/compile_fail/tile_export_outputs_split.rs"); +} 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 8c5b4a1da4..be20fed59d 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/api.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/api.rs @@ -667,6 +667,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()), @@ -689,6 +690,8 @@ impl Compiler { OutputFiles { llvm_ir: &ll_path, ptx: &ptx_path, + cubin: &cubin_path, + mlir: None, stale_before_export: &[], }, ); @@ -1011,6 +1014,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 f0cc44758a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 9a62dd9586..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 aa284eb54e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 a3e2c44a3b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 8955f8cc8c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 b50a23afae..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 f73fef865f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 0ada366dc9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 52c9ffdd8d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 735115d8a9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 dde0b8cefd..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 4f3de70be4..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 31cc8f6be6..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6c9be5ac8a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c6aa47c031..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 ea446d4825..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 74aa24dacd..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c7e5547922..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 1fd06948b5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `packed_conversion` intrinsics. @@ -718,4 +718,268 @@ pub(super) const TARGETS: &[GeneratedIntrinsicTarget] = &[ result_range: None, }), }, + GeneratedIntrinsicTarget { + marker: "v1:i1030", + id: "cvt_rn_satfinite_e4m3x2_bf16x2", + abi_id: "i1030", + dialect_op: "nvvm.cvt_rn_satfinite_e4m3x2_bf16x2", + variant: GeneratedIntrinsicVariant::Scalar, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + backend_requirements: &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + ], + selections: &[], + llvm: Some(GeneratedLlvmFacts { + properties: &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"], + result_no_undef: false, + result_range: None, + }), + }, + GeneratedIntrinsicTarget { + marker: "v1:i1031", + id: "cvt_rn_satfinite_relu_e4m3x2_bf16x2", + abi_id: "i1031", + dialect_op: "nvvm.cvt_rn_satfinite_relu_e4m3x2_bf16x2", + variant: GeneratedIntrinsicVariant::Scalar, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + backend_requirements: &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + ], + selections: &[], + llvm: Some(GeneratedLlvmFacts { + properties: &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"], + result_no_undef: false, + result_range: None, + }), + }, + GeneratedIntrinsicTarget { + marker: "v1:i1032", + id: "cvt_rn_satfinite_e5m2x2_bf16x2", + abi_id: "i1032", + dialect_op: "nvvm.cvt_rn_satfinite_e5m2x2_bf16x2", + variant: GeneratedIntrinsicVariant::Scalar, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + backend_requirements: &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + ], + selections: &[], + llvm: Some(GeneratedLlvmFacts { + properties: &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"], + result_no_undef: false, + result_range: None, + }), + }, + GeneratedIntrinsicTarget { + marker: "v1:i1033", + id: "cvt_rn_satfinite_relu_e5m2x2_bf16x2", + abi_id: "i1033", + dialect_op: "nvvm.cvt_rn_satfinite_relu_e5m2x2_bf16x2", + variant: GeneratedIntrinsicVariant::Scalar, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + backend_requirements: &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + ], + selections: &[], + llvm: Some(GeneratedLlvmFacts { + properties: &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"], + result_no_undef: false, + result_range: None, + }), + }, ]; 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 4c4a2f8287..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 bdceaba496..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 2ccb0c0db1..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `register_mma` intrinsics. @@ -7839,27 +7839,27 @@ pub(super) const TARGETS: &[GeneratedIntrinsicTarget] = &[ overflow: GeneratedRegisterMmaOverflow::NotApplicable, }, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm( - 80, + 75, )]), }, backend_requirements: &[ GeneratedBackendRequirement { backend: GeneratedIntrinsicBackend::LlvmNvptx, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[ - GeneratedHardwareAlternative::MinimumSm(80), + GeneratedHardwareAlternative::MinimumSm(75), ]), }, }, GeneratedBackendRequirement { backend: GeneratedIntrinsicBackend::LibNvvm, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[ - GeneratedHardwareAlternative::MinimumSm(80), + GeneratedHardwareAlternative::MinimumSm(75), ]), }, }, @@ -7959,27 +7959,27 @@ pub(super) const TARGETS: &[GeneratedIntrinsicTarget] = &[ overflow: GeneratedRegisterMmaOverflow::NotApplicable, }, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm( - 80, + 75, )]), }, backend_requirements: &[ GeneratedBackendRequirement { backend: GeneratedIntrinsicBackend::LlvmNvptx, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[ - GeneratedHardwareAlternative::MinimumSm(80), + GeneratedHardwareAlternative::MinimumSm(75), ]), }, }, GeneratedBackendRequirement { backend: GeneratedIntrinsicBackend::LibNvvm, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[ - GeneratedHardwareAlternative::MinimumSm(80), + GeneratedHardwareAlternative::MinimumSm(75), ]), }, }, @@ -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 102a59e264..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 66bef39204..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c0b78a50bf..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 712fbc5a18..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 8ac059d9cc..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 556a7e788e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c58ca91aab..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 e667da42b1..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 f9a55b392c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 ab96fe16aa..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 edd04c7507..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fa7086dcc6..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 a00308f306..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 2fd40c18ed..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6076281fa9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 cc0926521e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 4635787a76..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 4c2a20c02a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 123f15d932..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 8081db240d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 169baef8cf..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 78ab2245bb..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 faf89ff5ba..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 053f68fc88..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 675e1ca2cf..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 5e4809b884..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 b25414b0a5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 f34109e6e6..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 66b12afa5e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 75ff1a863c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 d8ceed3973..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6d7fad75c4..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `packed_conversion` intrinsics. @@ -802,3 +802,309 @@ fn cvt_rn_relu_f16x2_e5m2x2_target_matches_the_catalog() { &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"] as &[&str] ); } + +#[test] +fn cvt_rn_satfinite_e4m3x2_bf16x2_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1030").unwrap(); + assert_eq!(target.id, "cvt_rn_satfinite_e4m3x2_bf16x2"); + assert_eq!(target.abi_id, "i1030"); + assert_eq!(target.dialect_op, "nvvm.cvt_rn_satfinite_e4m3x2_bf16x2"); + assert_eq!(target.variant, GeneratedIntrinsicVariant::Scalar); + assert_eq!(target.requirement.minimum_ptx.encoded(), 91); + assert_eq!( + target.requirement.hardware, + const { + GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]) + } + ); + assert_eq!( + target.backend_requirements, + const { + &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + ] + } + ); + assert_eq!(target.selections, &[]); + assert_eq!( + target.llvm.unwrap().properties, + &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"] as &[&str] + ); +} + +#[test] +fn cvt_rn_satfinite_relu_e4m3x2_bf16x2_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1031").unwrap(); + assert_eq!(target.id, "cvt_rn_satfinite_relu_e4m3x2_bf16x2"); + assert_eq!(target.abi_id, "i1031"); + assert_eq!( + target.dialect_op, + "nvvm.cvt_rn_satfinite_relu_e4m3x2_bf16x2" + ); + assert_eq!(target.variant, GeneratedIntrinsicVariant::Scalar); + assert_eq!(target.requirement.minimum_ptx.encoded(), 91); + assert_eq!( + target.requirement.hardware, + const { + GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]) + } + ); + assert_eq!( + target.backend_requirements, + const { + &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + ] + } + ); + assert_eq!(target.selections, &[]); + assert_eq!( + target.llvm.unwrap().properties, + &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"] as &[&str] + ); +} + +#[test] +fn cvt_rn_satfinite_e5m2x2_bf16x2_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1032").unwrap(); + assert_eq!(target.id, "cvt_rn_satfinite_e5m2x2_bf16x2"); + assert_eq!(target.abi_id, "i1032"); + assert_eq!(target.dialect_op, "nvvm.cvt_rn_satfinite_e5m2x2_bf16x2"); + assert_eq!(target.variant, GeneratedIntrinsicVariant::Scalar); + assert_eq!(target.requirement.minimum_ptx.encoded(), 91); + assert_eq!( + target.requirement.hardware, + const { + GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]) + } + ); + assert_eq!( + target.backend_requirements, + const { + &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + ] + } + ); + assert_eq!(target.selections, &[]); + assert_eq!( + target.llvm.unwrap().properties, + &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"] as &[&str] + ); +} + +#[test] +fn cvt_rn_satfinite_relu_e5m2x2_bf16x2_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1033").unwrap(); + assert_eq!(target.id, "cvt_rn_satfinite_relu_e5m2x2_bf16x2"); + assert_eq!(target.abi_id, "i1033"); + assert_eq!( + target.dialect_op, + "nvvm.cvt_rn_satfinite_relu_e5m2x2_bf16x2" + ); + assert_eq!(target.variant, GeneratedIntrinsicVariant::Scalar); + assert_eq!(target.requirement.minimum_ptx.encoded(), 91); + assert_eq!( + target.requirement.hardware, + const { + GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]) + } + ); + assert_eq!( + target.backend_requirements, + const { + &[ + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LlvmNvptx, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + GeneratedBackendRequirement { + backend: GeneratedIntrinsicBackend::LibNvvm, + requirement: GeneratedTargetRequirement { + minimum_ptx: GeneratedPtxVersion::from_encoded(91), + hardware: GeneratedHardwareTarget::AnyOf(&[ + GeneratedHardwareAlternative::ExactArchitecture(100), + GeneratedHardwareAlternative::FamilyTarget(100), + GeneratedHardwareAlternative::ExactArchitecture(103), + GeneratedHardwareAlternative::FamilyTarget(103), + GeneratedHardwareAlternative::ExactArchitecture(110), + GeneratedHardwareAlternative::FamilyTarget(110), + GeneratedHardwareAlternative::ExactArchitecture(120), + GeneratedHardwareAlternative::FamilyTarget(120), + GeneratedHardwareAlternative::ExactArchitecture(121), + GeneratedHardwareAlternative::FamilyTarget(121), + ]), + }, + }, + ] + } + ); + assert_eq!(target.selections, &[]); + assert_eq!( + target.llvm.unwrap().properties, + &["IntrNoCreateUndefOrPoison", "IntrNoMem", "IntrSpeculatable"] as &[&str] + ); +} 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 214502fbda..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c4857e8d56..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 48b805f004..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -5572,10 +5572,10 @@ fn mma_m16n8k8_f16_f16_target_matches_the_catalog() { overflow: GeneratedRegisterMmaOverflow::NotApplicable } ); - assert_eq!(target.requirement.minimum_ptx.encoded(), 70); + assert_eq!(target.requirement.minimum_ptx.encoded(), 65); assert_eq!( target.requirement.hardware, - const { GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(80)]) } + const { GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(75)]) } ); assert_eq!( target.backend_requirements, @@ -5584,18 +5584,18 @@ fn mma_m16n8k8_f16_f16_target_matches_the_catalog() { GeneratedBackendRequirement { backend: GeneratedIntrinsicBackend::LlvmNvptx, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[ - GeneratedHardwareAlternative::MinimumSm(80), + GeneratedHardwareAlternative::MinimumSm(75), ]), }, }, GeneratedBackendRequirement { backend: GeneratedIntrinsicBackend::LibNvvm, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[ - GeneratedHardwareAlternative::MinimumSm(80), + GeneratedHardwareAlternative::MinimumSm(75), ]), }, }, @@ -5714,10 +5714,10 @@ fn mma_m16n8k8_f32_f16_target_matches_the_catalog() { overflow: GeneratedRegisterMmaOverflow::NotApplicable } ); - assert_eq!(target.requirement.minimum_ptx.encoded(), 70); + assert_eq!(target.requirement.minimum_ptx.encoded(), 65); assert_eq!( target.requirement.hardware, - const { GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(80)]) } + const { GeneratedHardwareTarget::AnyOf(&[GeneratedHardwareAlternative::MinimumSm(75)]) } ); assert_eq!( target.backend_requirements, @@ -5726,18 +5726,18 @@ fn mma_m16n8k8_f32_f16_target_matches_the_catalog() { GeneratedBackendRequirement { backend: GeneratedIntrinsicBackend::LlvmNvptx, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[ - GeneratedHardwareAlternative::MinimumSm(80), + GeneratedHardwareAlternative::MinimumSm(75), ]), }, }, GeneratedBackendRequirement { backend: GeneratedIntrinsicBackend::LibNvvm, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(70), + minimum_ptx: GeneratedPtxVersion::from_encoded(65), hardware: GeneratedHardwareTarget::AnyOf(&[ - GeneratedHardwareAlternative::MinimumSm(80), + GeneratedHardwareAlternative::MinimumSm(75), ]), }, }, @@ -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 5216dc659d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 ca62c3c7de..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 98af462f88..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 570b8b0232..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 bad07267a2..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 b800b7b44c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 aa3027531b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fc5ddff329..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 3dbb77ebd1..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 b01dcd7de4..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 90092d1fb8..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 589e374de8..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 e2ea3a9ddc..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 a324705f5e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 cf0d62e8d4..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target tests: `tma` intrinsics. @@ -6643,3 +6643,348 @@ fn cp_async_bulk_tensor_reduce_xor_im2col_5d_target_matches_the_catalog() { ] as &[&str] ); } + +#[test] +fn cp_async_bulk_tensor_1d_g2s_cta_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1034").unwrap(); + assert_eq!(target.id, "cp_async_bulk_tensor_1d_g2s_cta"); + assert_eq!(target.abi_id, "i1034"); + assert_eq!( + target.dialect_op, + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_1d" + ); + 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: "TMA_G2S_CTA_TILE_1D", + asm: "cp.async.bulk.tensor.1d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86" + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 5, + value: 0 + }] + } + },] + ); + assert_eq!( + target.llvm.unwrap().properties, + &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly" + ] as &[&str] + ); +} + +#[test] +fn cp_async_bulk_tensor_2d_g2s_cta_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1035").unwrap(); + assert_eq!(target.id, "cp_async_bulk_tensor_2d_g2s_cta"); + assert_eq!(target.abi_id, "i1035"); + assert_eq!( + target.dialect_op, + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_2d" + ); + 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: "TMA_G2S_CTA_TILE_2D", + asm: "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86" + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 6, + value: 0 + }] + } + },] + ); + assert_eq!( + target.llvm.unwrap().properties, + &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly" + ] as &[&str] + ); +} + +#[test] +fn cp_async_bulk_tensor_3d_g2s_cta_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1036").unwrap(); + assert_eq!(target.id, "cp_async_bulk_tensor_3d_g2s_cta"); + assert_eq!(target.abi_id, "i1036"); + assert_eq!( + target.dialect_op, + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_3d" + ); + 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: "TMA_G2S_CTA_TILE_3D", + asm: "cp.async.bulk.tensor.3d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86" + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 7, + value: 0 + }] + } + },] + ); + assert_eq!( + target.llvm.unwrap().properties, + &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly" + ] as &[&str] + ); +} + +#[test] +fn cp_async_bulk_tensor_4d_g2s_cta_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1037").unwrap(); + assert_eq!(target.id, "cp_async_bulk_tensor_4d_g2s_cta"); + assert_eq!(target.abi_id, "i1037"); + assert_eq!( + target.dialect_op, + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_4d" + ); + 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: "TMA_G2S_CTA_TILE_4D", + asm: "cp.async.bulk.tensor.4d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2, $d3}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86" + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 8, + value: 0 + }] + } + },] + ); + assert_eq!( + target.llvm.unwrap().properties, + &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly" + ] as &[&str] + ); +} + +#[test] +fn cp_async_bulk_tensor_5d_g2s_cta_target_matches_the_catalog() { + let target = generated_intrinsic_target_by_marker("v1:i1038").unwrap(); + assert_eq!(target.id, "cp_async_bulk_tensor_5d_g2s_cta"); + assert_eq!(target.abi_id, "i1038"); + assert_eq!( + target.dialect_op, + "nvvm.cp_async_bulk_tensor_g2s_cta_tile_5d" + ); + 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: "TMA_G2S_CTA_TILE_5D", + asm: "cp.async.bulk.tensor.5d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86" + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 9, + value: 0 + }] + } + },] + ); + assert_eq!( + target.llvm.unwrap().properties, + &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly" + ] as &[&str] + ); +} 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 34ab9c2e50..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 e0c0274b88..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 0a7b559a3b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 17608cd3ac..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 2fe8dd720f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c0323beb8d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated intrinsic target records: `tma` intrinsics. @@ -5670,4 +5670,294 @@ pub(super) const TARGETS: &[GeneratedIntrinsicTarget] = &[ result_range: None, }), }, + GeneratedIntrinsicTarget { + marker: "v1:i1034", + id: "cp_async_bulk_tensor_1d_g2s_cta", + abi_id: "i1034", + dialect_op: "nvvm.cp_async_bulk_tensor_g2s_cta_tile_1d", + 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: "TMA_G2S_CTA_TILE_1D", + asm: "cp.async.bulk.tensor.1d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86", + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 5, + value: 0, + }], + }, + }], + llvm: Some(GeneratedLlvmFacts { + properties: &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly", + ], + result_no_undef: false, + result_range: None, + }), + }, + GeneratedIntrinsicTarget { + marker: "v1:i1035", + id: "cp_async_bulk_tensor_2d_g2s_cta", + abi_id: "i1035", + dialect_op: "nvvm.cp_async_bulk_tensor_g2s_cta_tile_2d", + 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: "TMA_G2S_CTA_TILE_2D", + asm: "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86", + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 6, + value: 0, + }], + }, + }], + llvm: Some(GeneratedLlvmFacts { + properties: &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly", + ], + result_no_undef: false, + result_range: None, + }), + }, + GeneratedIntrinsicTarget { + marker: "v1:i1036", + id: "cp_async_bulk_tensor_3d_g2s_cta", + abi_id: "i1036", + dialect_op: "nvvm.cp_async_bulk_tensor_g2s_cta_tile_3d", + 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: "TMA_G2S_CTA_TILE_3D", + asm: "cp.async.bulk.tensor.3d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86", + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 7, + value: 0, + }], + }, + }], + llvm: Some(GeneratedLlvmFacts { + properties: &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly", + ], + result_no_undef: false, + result_range: None, + }), + }, + GeneratedIntrinsicTarget { + marker: "v1:i1037", + id: "cp_async_bulk_tensor_4d_g2s_cta", + abi_id: "i1037", + dialect_op: "nvvm.cp_async_bulk_tensor_g2s_cta_tile_4d", + 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: "TMA_G2S_CTA_TILE_4D", + asm: "cp.async.bulk.tensor.4d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2, $d3}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86", + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 8, + value: 0, + }], + }, + }], + llvm: Some(GeneratedLlvmFacts { + properties: &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly", + ], + result_no_undef: false, + result_range: None, + }), + }, + GeneratedIntrinsicTarget { + marker: "v1:i1038", + id: "cp_async_bulk_tensor_5d_g2s_cta", + abi_id: "i1038", + dialect_op: "nvvm.cp_async_bulk_tensor_g2s_cta_tile_5d", + 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: "TMA_G2S_CTA_TILE_5D", + asm: "cp.async.bulk.tensor.5d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$dst], [$tmap, {{$d0, $d1, $d2, $d3, $d4}}], [$mbar];", + predicates: &[ + "Subtarget->getSmVersion() >= 90", + "Subtarget->getPTXVersion() >= 86", + ], + constraints: GeneratedSelectionConstraints { + address_space: None, + immediate_bindings: &[GeneratedImmediateBinding { + argument_index: 9, + value: 0, + }], + }, + }], + llvm: Some(GeneratedLlvmFacts { + properties: &[ + "ImmArg", + "IntrConvergent", + "ReadOnly", + "WriteOnly", + ], + result_no_undef: false, + result_range: None, + }), + }, ]; 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 73d122c8b5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 b6829d7a5a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 8517fd196b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 20cb4a2df6..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 e06d161614..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 494e9e540a..1d4ac4ac87 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/iket.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/iket.rs @@ -37,7 +37,6 @@ use pliron::{ }, context::{Context, Ptr}, linked_list::ContainsLinkedList, - location::Located, op::{Op, op_cast}, operation::Operation, r#type::Typed, @@ -729,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( @@ -736,6 +737,8 @@ mod tests { nvvm, &backend, debug_kind, + None, + &[], files, PipelineTrace::default(), ); @@ -768,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, @@ -778,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/target/detect.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/target/detect.rs index 4e067ece0b..20ce82e87a 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/target/detect.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/target/detect.rs @@ -574,15 +574,17 @@ pub(super) fn contains_tma_cta_group_features(contents: &str) -> bool { }) } -/// Checks TMA copies whose destination is CTA-local shared memory. +/// Checks TMA copies whose destination is CTA-local shared memory, as inline +/// PTX or as the typed `llvm.nvvm.cp.async.bulk.tensor.g2s.cta.*` intrinsics. /// /// `.shared::cta` already existed as a source state space for shared-to-global /// copies, so the following `.global` source qualifier is part of the match. /// The destination form was introduced in PTX 8.6 but is valid on sm_90. pub(super) fn contains_tma_shared_cta_destination(contents: &str) -> bool { - contents.split(';').any(|statement| { - statement.contains("cp.async.bulk.") && statement.contains(".shared::cta.global") - }) + contents.contains("cp.async.bulk.tensor.g2s.cta.") + || contents.split(';').any(|statement| { + statement.contains("cp.async.bulk.") && statement.contains(".shared::cta.global") + }) } /// Checks PTX 8.6 TMA modifiers with a generic sm_100 architecture floor. diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/target/select.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/target/select.rs index cdc258e6b8..efee412907 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/target/select.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/target/select.rs @@ -252,18 +252,32 @@ pub fn validate_target_for_llvm_major( Ok(()) } +/// Refuses a PTX ISA requirement that the LLVM backend cannot honor. +/// +/// Each PTX feature the compiler can request has an LLVM floor: `+ptx90` needs +/// LLVM 22 and `+ptx91` needs LLVM 23. An older llc does not fail on a feature +/// it lacks; LLVM 22 ignores `+ptx91` with a warning and emits PTX 8.8, so the +/// build would pass and the module would fail to load. The gate refuses such a +/// backend at compile time instead. An unknown backend version cannot prove +/// support and is refused whenever a floor applies. pub(crate) fn validate_ptx_isa_for_llvm_major( requirement: PtxIsaRequirement, llc_major: Option, ) -> Result<(), String> { - if requirement >= PtxIsaRequirement::new(90) && llc_major.is_none_or(|major| major < 22) { - let backend = llc_major.map_or_else( - || "an LLVM backend with an unknown version".to_string(), - |major| format!("LLVM {major}"), - ); - return Err(format!( - "PTX 9.0 or newer requires LLVM 22 or newer; {backend} does not support the required PTX feature" - )); + // Newest floor first, so the message names the rule that applies. + const LLVM_FLOORS: [(u16, u32, &str); 2] = [(91, 23, "9.1"), (90, 22, "9.0")]; + for (spelling, llvm_major, label) in LLVM_FLOORS { + if requirement >= PtxIsaRequirement::new(spelling) + && llc_major.is_none_or(|major| major < llvm_major) + { + let backend = llc_major.map_or_else( + || "an LLVM backend with an unknown version".to_string(), + |major| format!("LLVM {major}"), + ); + return Err(format!( + "PTX {label} or newer requires LLVM {llvm_major} or newer; {backend} does not support the required PTX feature" + )); + } } Ok(()) } diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/detection_async_memory.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/detection_async_memory.rs index a81d615885..9c644810bb 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/detection_async_memory.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/detection_async_memory.rs @@ -58,6 +58,15 @@ fn tma_and_wgmma_raise_their_independent_ptx_floors() { assert_eq!(shared_cta_requirements.features, DetectedFeatures::Tma); assert_eq!(shared_cta_requirements.ptx_isa, PtxIsaRequirement::new(86)); + let typed_shared_cta = "call void @llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.2d(ptr addrspace(3) %0, ptr addrspace(3) %1, ptr %2, i32 %3, i32 %4, i64 0, i1 false)"; + assert!(contains_tma_shared_cta_destination(typed_shared_cta)); + let typed_requirements = detect_module_requirements_in_llvm_text(typed_shared_cta); + assert_eq!(typed_requirements.features, DetectedFeatures::Tma); + assert_eq!(typed_requirements.ptx_isa, PtxIsaRequirement::new(86)); + assert!(!contains_tma_shared_cta_destination( + "call void @llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d(ptr addrspace(7) %0)" + )); + let shared_source = "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group;"; assert!(!contains_tma_shared_cta_destination(shared_source)); assert_eq!( diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/fixtures.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/fixtures.rs index e06a2abce7..17020cc268 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/fixtures.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/fixtures.rs @@ -54,14 +54,14 @@ pub(super) static PTX90: GeneratedIntrinsicTarget = GeneratedIntrinsicTarget { selections: &[], llvm: None, }; -pub(super) static PTX91_FUTURE: GeneratedIntrinsicTarget = GeneratedIntrinsicTarget { - marker: "test:ptx91", - id: "ptx91_future", +pub(super) static PTX92_FUTURE: GeneratedIntrinsicTarget = GeneratedIntrinsicTarget { + marker: "test:ptx92", + id: "ptx92_future", abi_id: "test", - dialect_op: "test.ptx91", + dialect_op: "test.ptx92", variant: GeneratedIntrinsicVariant::Scalar, requirement: GeneratedTargetRequirement { - minimum_ptx: GeneratedPtxVersion::from_encoded(91), + minimum_ptx: GeneratedPtxVersion::from_encoded(92), hardware: GeneratedHardwareTarget::All, }, backend_requirements: &[], diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/generated_floors.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/generated_floors.rs index 52ed6075ce..1015e965a3 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/generated_floors.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/target/tests/generated_floors.rs @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -use super::fixtures::{PTX87_EXACT_SM120A, PTX88, PTX90, PTX91_FUTURE, TCGEN_F16, TCGEN_I8}; +use super::fixtures::{PTX87_EXACT_SM120A, PTX88, PTX90, PTX92_FUTURE, TCGEN_F16, TCGEN_I8}; use crate::generated::GeneratedModuleRequirements; use crate::generated_intrinsic_targets::{ GeneratedHardwareAlternative, GeneratedHardwareTarget, GeneratedIntrinsicBackend, @@ -482,10 +482,28 @@ fn generated_ptx88_and_ptx90_floors_are_preserved() { assert!(validate_ptx_isa_for_llvm_major(PtxIsaRequirement::new(90), Some(21)).is_err()); validate_ptx_isa_for_llvm_major(PtxIsaRequirement::new(90), Some(22)).unwrap(); - let generated = GeneratedModuleRequirements::from_targets(vec![&PTX91_FUTURE]); + // PTX 9.1 is requestable and needs LLVM 23. LLVM 22 ignores `+ptx91` and + // emits PTX 8.8, so the gate has to refuse it at compile time. + assert_eq!( + ptx_isa_requirement_for_floor(91, "test", "test").unwrap(), + PtxIsaRequirement::new(91) + ); + assert_eq!( + required_ptx_feature(&"sm_120f".parse().unwrap(), PtxIsaRequirement::new(91)).unwrap(), + Some("+ptx91") + ); + assert!(validate_ptx_isa_for_llvm_major(PtxIsaRequirement::new(91), None).is_err()); + let error = validate_ptx_isa_for_llvm_major(PtxIsaRequirement::new(91), Some(22)).unwrap_err(); + assert!( + error.contains("PTX 9.1 or newer requires LLVM 23 or newer"), + "{error}" + ); + validate_ptx_isa_for_llvm_major(PtxIsaRequirement::new(91), Some(23)).unwrap(); + + let generated = GeneratedModuleRequirements::from_targets(vec![&PTX92_FUTURE]); let error = generated_ptx_isa_requirement(&generated).unwrap_err(); - assert!(error.contains("requires PTX 9.1"), "{error}"); + assert!(error.contains("requires PTX 9.2"), "{error}"); assert!( error.contains("newer than cuda-oxide can request"), "{error}" @@ -750,3 +768,67 @@ fn generated_dot_product_floors_record_sm61_and_split_backend_support() { )); } } + +#[test] +fn generated_m16n8k8_mma_floors_admit_turing_only_for_fp16() { + use crate::generated_intrinsic_targets::{ + GeneratedIntrinsicBackend, generated_intrinsic_target_by_marker, + }; + + let backends = [ + GeneratedIntrinsicBackend::LlvmNvptx, + GeneratedIntrinsicBackend::LibNvvm, + ]; + // mma_m16n8k8_f16_f16 and mma_m16n8k8_f32_f16 + for marker in ["v1:i0521", "v1:i0523"] { + let target = generated_intrinsic_target_by_marker(marker).unwrap(); + for backend in backends { + let generated = + GeneratedModuleRequirements::from_targets(vec![target]).for_backend(backend); + assert_eq!( + generated_ptx_isa_requirement(&generated).unwrap(), + PtxIsaRequirement::new(65), + "{marker} {backend:?}" + ); + assert!( + !generated_target_satisfied(&"sm_72".parse().unwrap(), &generated), + "{marker} {backend:?}" + ); + for arch in ["sm_75", "sm_80"] { + assert!( + generated_target_satisfied(&arch.parse().unwrap(), &generated), + "{marker} {backend:?} {arch}" + ); + } + validate_generated_target(&"sm_75".parse().unwrap(), &generated).unwrap(); + } + } + + // The three sm_80 members of the admission (m16n8k4 TF32, m16n8k8 BF16, + // m16n8k16 F16) and the m16n8k8 TF32 form keep the Ampere floor, although + // LLVM selects the BF16 and TF32 m16n8k8 forms at sm_75. + for marker in ["v1:i0520", "v1:i0522", "v1:i0524", "v1:i0107"] { + let target = generated_intrinsic_target_by_marker(marker).unwrap(); + for backend in backends { + let generated = + GeneratedModuleRequirements::from_targets(vec![target]).for_backend(backend); + assert_eq!( + generated_ptx_isa_requirement(&generated).unwrap(), + PtxIsaRequirement::new(70), + "{marker} {backend:?}" + ); + assert!( + !generated_target_satisfied(&"sm_75".parse().unwrap(), &generated), + "{marker} {backend:?}" + ); + assert!( + generated_target_satisfied(&"sm_80".parse().unwrap(), &generated), + "{marker} {backend:?}" + ); + let error = + validate_generated_target(&"sm_75".parse().unwrap(), &generated).unwrap_err(); + assert!(error.contains(target.id), "{error}"); + assert!(error.contains("sm_80 or newer"), "{error}"); + } + } +} 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/cuda-target-spec/src/lib.rs b/cuda-oxide/crates/cuda-target-spec/src/lib.rs index a1ef76a680..62b2f65d99 100644 --- a/cuda-oxide/crates/cuda-target-spec/src/lib.rs +++ b/cuda-oxide/crates/cuda-target-spec/src/lib.rs @@ -481,7 +481,7 @@ pub fn recorded_ptx_floor(arch: &CudaArch) -> Result "+ptx87", 88 => "+ptx88", 90 => "+ptx90", + 91 => "+ptx91", _ => unreachable!(), } } @@ -670,7 +671,8 @@ mod tests { assert_eq!(spelling_feature(74), None); assert_eq!(PtxSpelling::from_spelling(74), None); assert_eq!(PtxSpelling::round_up(74).map(PtxSpelling::get), Some(78)); - assert_eq!(PtxSpelling::round_up(91), None); + assert_eq!(PtxSpelling::round_up(91).map(PtxSpelling::get), Some(91)); + assert_eq!(PtxSpelling::round_up(92), None); } #[test] diff --git a/cuda-oxide/crates/dialect-iket/src/ops.rs b/cuda-oxide/crates/dialect-iket/src/ops.rs index 8ce56c26bd..1ac81c6ddf 100644 --- a/cuda-oxide/crates/dialect-iket/src/ops.rs +++ b/cuda-oxide/crates/dialect-iket/src/ops.rs @@ -14,7 +14,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/dialect-mir/src/attributes.rs b/cuda-oxide/crates/dialect-mir/src/attributes.rs index 1152f531e3..cb287329f4 100644 --- a/cuda-oxide/crates/dialect-mir/src/attributes.rs +++ b/cuda-oxide/crates/dialect-mir/src/attributes.rs @@ -79,19 +79,51 @@ pub struct MutabilityAttr(pub bool); /// Proven validity facts for one Rust-reference kernel parameter. /// /// Presence of this attribute means the source argument is a Rust reference -/// proven non-null by `rustc_public`. The payload is the pointee alignment in -/// bytes; it must be a non-zero power of two. LLVM `align 1` is redundant and +/// proven non-null by `rustc_public`. Pointee alignment is in bytes and must +/// be a non-zero power of two. LLVM `align 1` is redundant and /// is omitted during export, but retaining `1` here keeps the proof complete. /// /// The attribute is stored on `mir.func` under a source-argument-indexed key. -/// It intentionally carries no aliasing or dereferenceability promise. -#[pliron_attr( - name = "mir.reference_param_validity", - format = "$0", - verifier = "succ" -)] +/// Additional promises are independently proven by the importer. A zero +/// dereferenceable byte count records no dereferenceability promise. +#[pliron_attr(name = "mir.reference_param_validity", format, verifier = "succ")] #[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] -pub struct ReferenceParamValidityAttr(pub u64); +pub struct ReferenceParamValidityAttr { + pub pointee_alignment: u64, + pub no_alias: bool, + pub dereferenceable_bytes: u64, +} + +impl ReferenceParamValidityAttr { + /// Non-nullness and alignment, without additional reference promises. + pub const fn new(pointee_alignment: u64) -> Self { + Self { + pointee_alignment, + no_alias: false, + dereferenceable_bytes: 0, + } + } +} + +/// The inline request a source function carries, recorded on `mir.func`. +/// +/// `#[inline]` is `Hint`, `#[inline(always)]` is `Always` and +/// `#[inline(never)]` is `Never`. A function without the attribute carries +/// none, and a kernel entry never carries one: it is never called, so there +/// is nothing to inline. The dialect records only the request; `mir-lower` +/// chooses the LLVM function attribute that expresses it. +#[pliron_attr(name = "mir.inline_intent", format, verifier = "succ")] +#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)] +pub enum InlineIntentAttr { + Hint, + Always, + Never, +} + +impl InlineIntentAttr { + /// Every request, for tests that cover each one. + pub const ALL: [Self; 3] = [Self::Hint, Self::Always, Self::Never]; +} /// Structural field index for aggregate access ops /// (`mir.extract_field`, `mir.insert_field`, `mir.field_addr`, `mir.enum_payload`). @@ -200,6 +232,7 @@ pub fn register(ctx: &mut Context) { MirPointerKindAuthorityAttr::register(ctx); MutabilityAttr::register(ctx); ReferenceParamValidityAttr::register(ctx); + InlineIntentAttr::register(ctx); FieldIndexAttr::register(ctx); VariantIndexAttr::register(ctx); UnrollAttr::register(ctx); diff --git a/cuda-oxide/crates/dialect-mir/src/ops/aggregate.rs b/cuda-oxide/crates/dialect-mir/src/ops/aggregate.rs index ceb6305021..0fd6415c69 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/aggregate.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/aggregate.rs @@ -14,7 +14,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, printable::Printable, diff --git a/cuda-oxide/crates/dialect-mir/src/ops/arithmetic.rs b/cuda-oxide/crates/dialect-mir/src/ops/arithmetic.rs index 7f0e34ab13..3308f367b1 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/arithmetic.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/arithmetic.rs @@ -14,7 +14,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/dialect-mir/src/ops/call.rs b/cuda-oxide/crates/dialect-mir/src/ops/call.rs index f40deb5d26..8fde9fe530 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/call.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/call.rs @@ -25,7 +25,6 @@ use pliron::{ common_traits::Verify, context::{Context, Ptr}, linked_list::ContainsLinkedList, - location::Located, op::Op, operation::Operation, result::Error, @@ -92,6 +91,38 @@ impl MirCallOp { self.set_attr_external_callee_type(ctx, TypeAttr::new(signature)); self.set_attr_call_pointer_kind_authority(ctx, MirPointerKindAuthorityAttr::AbiBoundary); } + + /// Whether the call carries Rust's logical `()` result for a void callee. + /// Explicitly declared aggregate results, including empty tuples, are not + /// this bridge and must retain their ABI result when exported. + pub fn has_logical_unit_result_bridge(&self, ctx: &Context) -> bool { + let op = self.get_operation().deref(ctx); + if op.get_num_results() != 1 + || !op + .get_result(0) + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .is_some_and(|tuple| tuple.get_types().is_empty()) + { + return false; + } + let Some(name) = self.get_attr_callee(ctx) else { + return false; + }; + let signature = match resolve_call(ctx, self.get_operation(), name.as_str()) { + CallResolution::AttachedResolved(callee) => Some(callee.get_type(ctx).into()), + _ => self.get_attr_external_callee_type(ctx).map(|attr| { + attr_cast::(&*attr) + .expect("TypeAttr must implement TypedAttrInterface") + .get_type(ctx) + }), + }; + signature.is_some_and(|ty| { + type_cast::(&*ty.deref(ctx)) + .is_some_and(|function| function.res_types().is_empty()) + }) + } } impl Verify for MirCallOp { diff --git a/cuda-oxide/crates/dialect-mir/src/ops/cast.rs b/cuda-oxide/crates/dialect-mir/src/ops/cast.rs index fe796219d8..500d066862 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/cast.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/cast.rs @@ -22,7 +22,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/dialect-mir/src/ops/comparison.rs b/cuda-oxide/crates/dialect-mir/src/ops/comparison.rs index 72f9cab4f1..080aa3f6f3 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/comparison.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/comparison.rs @@ -14,7 +14,6 @@ use pliron::{ }, common_traits::Verify, context::Context, - location::Located, op::Op, printable::Printable, result::Error, diff --git a/cuda-oxide/crates/dialect-mir/src/ops/constants.rs b/cuda-oxide/crates/dialect-mir/src/ops/constants.rs index 615c5c2ff2..63b6e3398e 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/constants.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/constants.rs @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs b/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs index 358f457b1c..fc2d7f51ed 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/control_flow.rs @@ -7,10 +7,14 @@ //! //! This module defines assertions, terminators, and branches for the MIR dialect. +use std::ops::Range; + use pliron::{ builtin::{ + attributes::OperandSegmentSizesAttr, op_interfaces::{ - BranchOpInterface, IsTerminatorInterface, NResultsInterface, OperandSegmentInterface, + ATTR_KEY_OPERAND_SEGMENT_SIZES, BranchOpInterface, IsTerminatorInterface, + NResultsInterface, OperandSegmentInterface, }, type_interfaces::FunctionTypeInterface, types::IntegerType, @@ -18,7 +22,6 @@ use pliron::{ common_traits::Verify, context::{Context, Ptr}, derive::op_interface_impl, - location::Located, op::Op, operation::Operation, result::Error, @@ -146,41 +149,23 @@ impl MirGotoOp { impl Verify for MirGotoOp { fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = &*self.get_operation().deref(ctx); + ::verify(self, ctx) + } +} +#[op_interface_impl] +impl BranchOpInterface for MirGotoOp { + fn verify_successor_operand_layout(&self, ctx: &Context) -> Result<(), Error> { + let op = self.get_operation().deref(ctx); if op.get_num_successors() != 1 { return verify_err!(op.loc(), "MirGotoOp must have exactly 1 successor"); } - let succ = op.get_successor(0); - let succ_block = succ.deref(ctx); - - if op.get_num_operands() != succ_block.get_num_arguments() { - return verify_err!( - op.loc(), - "MirGotoOp operand count must match successor argument count" - ); - } - - for i in 0..succ_block.get_num_arguments() { - let arg = succ_block.get_argument(i); - let opd = op.get_operand(i); - if opd.get_type(ctx) != arg.get_type(ctx) { - return verify_err!( - op.loc(), - "MirGotoOp operand type mismatch with successor argument" - ); - } - } - Ok(()) } -} -#[op_interface_impl] -impl BranchOpInterface for MirGotoOp { - fn successor_operands(&self, ctx: &Context, succ_idx: usize) -> Vec { + fn successor_operand_range(&self, ctx: &Context, succ_idx: usize) -> Range { assert!(succ_idx == 0, "MirGotoOp has exactly one successor"); - self.get_operation().deref(ctx).operands().collect() + 0..self.get_operation().deref(ctx).get_num_operands() } fn add_successor_operand(&self, ctx: &mut Context, succ_idx: usize, operand: Value) -> usize { @@ -218,13 +203,14 @@ impl BranchOpInterface for MirGotoOp { /// /// # Verification /// +/// - Exactly three operand segments, with one condition operand. /// - Condition must be `i1`. /// - Must have exactly 2 successors. /// - Operand types must match successor block argument types. #[pliron_op( name = "mir.cond_br", format, - interfaces = [NResultsInterface<0>, IsTerminatorInterface, OperandSegmentInterface] + interfaces = [NResultsInterface<0>, IsTerminatorInterface] )] pub struct MirCondBranchOp; @@ -237,20 +223,10 @@ impl MirCondBranchOp { impl Verify for MirCondBranchOp { fn verify(&self, ctx: &Context) -> Result<(), Error> { - let op = &*self.get_operation().deref(ctx); - - if op.get_num_successors() != 2 { - return verify_err!(op.loc(), "MirCondBranchOp must have exactly 2 successors"); - } - let true_block = op.get_successor(0).deref(ctx); - let false_block = op.get_successor(1).deref(ctx); - - if op.get_num_operands() < 1 { - return verify_err!( - op.loc(), - "MirCondBranchOp must have at least 1 operand (condition)" - ); - } + // Validate segment metadata before reading operands. The shared branch + // verifier also checks each segment against its target block arguments. + ::verify(self, ctx)?; + let op = self.get_operation().deref(ctx); let cond = op.get_operand(0); let cond_ty = cond.get_type(ctx); let cond_ty_obj = cond_ty.deref(ctx); @@ -263,55 +239,83 @@ impl Verify for MirCondBranchOp { return verify_err!(op.loc(), "MirCondBranchOp condition must be integer type"); } - let expected_args = 1 + true_block.get_num_arguments() + false_block.get_num_arguments(); - if op.get_num_operands() != expected_args { + Ok(()) + } +} + +#[op_interface_impl] +impl OperandSegmentInterface for MirCondBranchOp { + fn expected_num_segments(&self, _ctx: &Context) -> Option { + // Condition, true successor arguments, false successor arguments. + Some(3) + } + + fn verify(op: &dyn Op, ctx: &Context) -> Result<(), Error> { + let op = op.get_operation().deref(ctx); + let Some(sizes) = op + .attributes + .get::(&ATTR_KEY_OPERAND_SEGMENT_SIZES) + else { + return verify_err!(op.loc(), "MirCondBranchOp requires operand segment sizes"); + }; + if sizes.0.len() != 3 { return verify_err!( op.loc(), - "MirCondBranchOp operand count must match condition + successors arguments" + "MirCondBranchOp must have exactly 3 operand segments" ); } - - let mut op_idx = 1; - for i in 0..true_block.get_num_arguments() { - let arg = true_block.get_argument(i); - let opd = op.get_operand(op_idx); - if opd.get_type(ctx) != arg.get_type(ctx) { - return verify_err!( - op.loc(), - "MirCondBranchOp true successor argument type mismatch" - ); - } - op_idx += 1; - } - - for i in 0..false_block.get_num_arguments() { - let arg = false_block.get_argument(i); - let opd = op.get_operand(op_idx); - if opd.get_type(ctx) != arg.get_type(ctx) { - return verify_err!( - op.loc(), - "MirCondBranchOp false successor argument type mismatch" - ); - } - op_idx += 1; + // Pliron 0.19's default verifier sums in u32 without checking overflow. + // Keep its range arithmetic valid even for untrusted segment metadata. + let Some(total) = sizes + .0 + .iter() + .try_fold(0u32, |sum, size| sum.checked_add(*size)) + else { + return verify_err!( + op.loc(), + "MirCondBranchOp operand segment total overflows u32" + ); + }; + if usize::try_from(total).ok() != Some(op.get_num_operands()) { + return verify_err!( + op.loc(), + "MirCondBranchOp operand segment total must match operand count" + ); } - Ok(()) } } #[op_interface_impl] impl BranchOpInterface for MirCondBranchOp { - fn successor_operands(&self, ctx: &Context, succ_idx: usize) -> Vec { + fn verify_successor_operand_layout(&self, ctx: &Context) -> Result<(), Error> { + // Interface verification can run before the operation's own verifier. + // Check every precondition of segment_range here, before it is called. + ::verify(self, ctx)?; + let op = self.get_operation().deref(ctx); + if op.get_num_successors() != 2 { + return verify_err!(op.loc(), "MirCondBranchOp must have exactly 2 successors"); + } + if self.segment_size(ctx, 0) != 1 { + return verify_err!( + op.loc(), + "MirCondBranchOp must have exactly 1 condition operand" + ); + } + Ok(()) + } + + fn successor_operand_range(&self, ctx: &Context, succ_idx: usize) -> Range { assert!( succ_idx == 0 || succ_idx == 1, "MirCondBranchOp has exactly two successors" ); - // Segment 0 = condition, segment 1 = true args, segment 2 = false args - self.get_segment(ctx, succ_idx + 1) + // Segment 0 = condition, segment 1 = true args, segment 2 = false args. + self.segment_range(ctx, succ_idx + 1) } fn add_successor_operand(&self, ctx: &mut Context, succ_idx: usize, operand: Value) -> usize { + assert!(succ_idx < 2, "MirCondBranchOp has exactly two successors"); self.push_to_segment(ctx, succ_idx + 1, operand) } @@ -321,6 +325,7 @@ impl BranchOpInterface for MirCondBranchOp { succ_idx: usize, opd_idx: usize, ) -> Value { + assert!(succ_idx < 2, "MirCondBranchOp has exactly two successors"); self.remove_from_segment(ctx, succ_idx + 1, opd_idx) } } diff --git a/cuda-oxide/crates/dialect-mir/src/ops/debug.rs b/cuda-oxide/crates/dialect-mir/src/ops/debug.rs index 89e8251a05..a64e8cd6dc 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/debug.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/debug.rs @@ -17,7 +17,7 @@ use pliron::{ common_traits::Verify, context::{Context, Ptr}, identifier::Identifier, - location::{Located, Location}, + location::Location, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/dialect-mir/src/ops/enum_ops.rs b/cuda-oxide/crates/dialect-mir/src/ops/enum_ops.rs index d7d9989b89..b5200e947c 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/enum_ops.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/enum_ops.rs @@ -13,7 +13,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, printable::Printable, diff --git a/cuda-oxide/crates/dialect-mir/src/ops/function.rs b/cuda-oxide/crates/dialect-mir/src/ops/function.rs index 8ed3037336..cb1deda4c9 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/function.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/function.rs @@ -25,16 +25,15 @@ use pliron::{ common_traits::Verify, context::{Context, Ptr}, identifier::Identifier, - indented_block, input_err, + input_err, irfmt::{ parsers::{spaced, type_parser}, printers::op::{region, typed_symb_op_header}, }, linked_list::ContainsLinkedList, - location::Located, op::{Op, OpObj}, operation::Operation, - parsable::{Parsable, ParseResult, StateStream}, + parsable::{Parsable, ParseResult, StateStream, StreamLocation}, printable::{Printable, State, indented_nl}, region::Region, result::Error, @@ -44,17 +43,25 @@ use pliron::{ use pliron_derive::pliron_op; use crate::{ - attributes::ReferenceParamValidityAttr, + attributes::{InlineIntentAttr, ReferenceParamValidityAttr}, types::{MirPtrType, MirSliceType}, }; const REFERENCE_PARAM_VALIDITY_ATTR_PREFIX: &str = "reference_param_validity_"; +/// Operation attribute key carrying a `mir.func`'s [`InlineIntentAttr`]. +pub const INLINE_INTENT_ATTR_KEY: &str = "inline_intent"; + fn reference_param_validity_key(index: usize) -> Identifier { Identifier::try_new(format!("{REFERENCE_PARAM_VALIDITY_ATTR_PREFIX}{index}")) .expect("reference parameter validity attribute name is valid") } +fn inline_intent_key() -> Identifier { + Identifier::try_new(INLINE_INTENT_ATTR_KEY.to_string()) + .expect("inline intent attribute name is valid") +} + /// MIR function operation. /// /// Represents a function in MIR. Contains a single region with basic blocks. @@ -66,7 +73,8 @@ fn reference_param_validity_key(index: usize) -> Identifier { /// |----------------|-----------|------------------------------------| /// | `sym_name` | StringAttr| Function name (from SymbolOpInterface) | /// | `mir_func_type`| TypeAttr | Function type (mir.func_type) | -/// | `reference_param_validity_N` | ReferenceParamValidityAttr | Proven nonnull/alignment for source argument `N` on a kernel entry | +/// | `reference_param_validity_N` | ReferenceParamValidityAttr | Proven reference facts for source argument `N` on a kernel entry | +/// | `inline_intent` | InlineIntentAttr | The source function's `#[inline]` request; never on a kernel entry | /// ``` /// /// # Verification @@ -116,8 +124,9 @@ impl MirFuncOp { /// Record rustc-proven validity for one source-level kernel argument. /// - /// Presence proves `nonnull`; the payload carries the rustc ABI alignment - /// of the pointee. The source argument index is deliberately retained until + /// Presence proves `nonnull`; the payload carries the pointee alignment + /// and independently proven aliasing and size facts. + /// The source argument index is deliberately retained until /// kernel ABI lowering decides which physical LLVM parameter represents it. pub fn set_reference_param_validity( &self, @@ -143,6 +152,24 @@ impl MirFuncOp { .get::(&reference_param_validity_key(index)) .copied() } + + /// Record the source function's `#[inline]` request. `mir-lower` turns it + /// into the LLVM function attribute; kernels must carry none (verified). + pub fn set_inline_intent(&self, ctx: &mut Context, intent: InlineIntentAttr) { + self.get_operation() + .deref_mut(ctx) + .attributes + .set(inline_intent_key(), intent); + } + + /// The recorded `#[inline]` request; `None` when the source had none. + pub fn inline_intent(&self, ctx: &Context) -> Option { + self.get_operation() + .deref(ctx) + .attributes + .get::(&inline_intent_key()) + .copied() + } } impl Typed for MirFuncOp { @@ -170,10 +197,9 @@ impl Printable for MirFuncOp { .retain(|key, _| key != &*ATTR_KEY_MIR_FUNC_TYPE && key != &ATTR_KEY_SYM_NAME); if !attributes_to_print_separately.0.is_empty() { - indented_block!(state, { - write!(f, "{}", indented_nl(state))?; - attributes_to_print_separately.fmt(ctx, state, f)?; - }); + let _indent = state.indent(); + write!(f, "{}", indented_nl(state))?; + attributes_to_print_separately.fmt(ctx, state, f)?; } write!(f, " ")?; region(self).fmt(ctx, state, f)?; @@ -201,19 +227,25 @@ impl Parsable for MirFuncOp { vec![], vec![], vec![], - 1, + // Region::parser creates the body region itself. + 0, ); let mut parser = ( spaced(token('@').with(Identifier::parser(()))).skip(spaced(token(':'))), spaced(type_parser()), - spaced(AttributeDict::parser(())), + spaced(optional(AttributeDict::parser(()))), spaced(optional(Region::parser(op))), ); parser .parse_stream(state_stream) - .map(|(fname, fty, attrs, _region)| -> OpObj { + .map(|(fname, fty, attrs, region)| -> OpObj { let ctx = &mut state_stream.state.ctx; - op.deref_mut(ctx).attributes = attrs; + // Declarations still have one empty region, matching the + // function operation's OneRegionInterface contract. + if region.is_none() { + Operation::add_region(op, ctx); + } + op.deref_mut(ctx).attributes = attrs.unwrap_or_default(); let ty_attr = TypeAttr::new(fty); let opop = MirFuncOp { op }; opop.set_symbol_name(ctx, fname); @@ -250,6 +282,31 @@ impl Verify for MirFuncOp { let is_kernel = op.attributes.get::(&kernel_key).is_some(); let inputs = interface.arg_types(); + // An inline request only applies to callees. A kernel `.entry` is + // never called, so there is no call to inline and a request on it + // means nothing; LLVM itself accepts the keywords on a `ptx_kernel`. + // The key must also hold the dialect's own attribute type, so a + // request can never be read back as "none" by mistake. + let inline_intent_key = inline_intent_key(); + if op.attributes.0.contains_key(&inline_intent_key) + && op + .attributes + .get::(&inline_intent_key) + .is_none() + { + return verify_err!( + op.loc(), + "MirFuncOp inline intent attribute has the wrong attribute type" + ); + } + if is_kernel && let Some(intent) = self.inline_intent(ctx) { + return verify_err!( + op.loc(), + "MirFuncOp inline intent {:?} is not valid on a kernel entry", + intent + ); + } + for (key, _) in &op.attributes.0 { let key_text = key.to_string(); let Some(index_text) = key_text.strip_prefix(REFERENCE_PARAM_VALIDITY_ATTR_PREFIX) @@ -285,12 +342,12 @@ impl Verify for MirFuncOp { key_text ); }; - if validity.0 == 0 || !validity.0.is_power_of_two() { + if validity.pointee_alignment == 0 || !validity.pointee_alignment.is_power_of_two() { return verify_err!( op.loc(), "MirFuncOp reference validity alignment on argument {} must be a non-zero power of two, found {}", index, - validity.0 + validity.pointee_alignment ); } diff --git a/cuda-oxide/crates/dialect-mir/src/ops/memory.rs b/cuda-oxide/crates/dialect-mir/src/ops/memory.rs index b59023dbf9..285c7b9723 100644 --- a/cuda-oxide/crates/dialect-mir/src/ops/memory.rs +++ b/cuda-oxide/crates/dialect-mir/src/ops/memory.rs @@ -18,7 +18,6 @@ use pliron::{ derive::op_interface_impl, identifier::Identifier, irbuild::{inserter::Inserter, rewriter::Rewriter}, - location::Located, op::Op, operation::Operation, opts::mem2reg::{ @@ -1117,6 +1116,16 @@ impl Verify for MirSharedAllocOp { return verify_err!(op.loc(), "MirSharedAllocOp missing size attribute"); } + if let Some(alignment) = self.get_alignment_value(ctx) + && (alignment == 0 || !alignment.is_power_of_two()) + { + return verify_err!( + op.loc(), + "MirSharedAllocOp alignment must be a non-zero power of two, found {}", + alignment + ); + } + // Check result is a shared memory pointer let res = op.get_result(0); let res_ty = res.get_type(ctx); @@ -1354,18 +1363,21 @@ mod device_global_key_tests { /// All `MirExternSharedOp` instances in a kernel refer to the **same** underlying /// memory, with different byte offsets for partitioning. /// -/// This is lowered to an LLVM global with extern linkage and addrspace(3): +/// This is lowered to an LLVM global with extern linkage and addrspace(3). +/// The global alignment is the maximum effective requirement across the +/// function's dynamic shared-memory references and any propagated launch +/// contract: /// ```llvm -/// @__dynamic_smem = external addrspace(3) global [0 x i8], align 256 +/// @__dynamic_smem = external addrspace(3) global [0 x i8], align /// ``` /// /// # Attributes /// /// ```text -/// | Name | Type | Description | -/// |-----------------|-------------|------------------------------------------| -/// | `byte_offset` | IntegerAttr | Byte offset from start of dynamic smem | -/// | `mir_alignment` | IntegerAttr | Alignment hint (global always uses 256) | +/// | Name | Type | Description | +/// |----------------------|-------------|------------------------------------------------| +/// | `extern_byte_offset` | IntegerAttr | Byte offset from start of dynamic smem | +/// | `extern_alignment` | IntegerAttr | Requested minimum alignment for this reference | /// ``` /// /// # Results @@ -1379,6 +1391,10 @@ mod device_global_key_tests { /// # Verification /// /// - Result must be a pointer type with shared address space (3). +/// - An explicit alignment must be a non-zero power of two. +/// +/// Lowering also raises the requested alignment to the pointee type's required +/// ABI alignment before computing the maximum alignment for the shared pool. #[pliron_op( name = "mir.extern_shared", format, @@ -1414,10 +1430,11 @@ impl MirExternSharedOp { self.set_attr_extern_byte_offset(ctx, offset_attr); } - /// Get alignment as u64 (returns 128 if not set). + /// Get the requested minimum alignment as u64. /// - /// Note: The actual global alignment is fixed at 256 bytes in mir-lower, - /// regardless of this attribute value. + /// Legacy or synthetic MIR without an explicit alignment attribute falls + /// back to 128 bytes. Importer-produced dynamic shared-memory operations + /// carry the source-level alignment explicitly. pub fn get_alignment_value(&self, ctx: &Context) -> u64 { self.get_attr_extern_alignment(ctx) .map(|attr| attr.value().to_u64()) @@ -1443,6 +1460,18 @@ impl Verify for MirExternSharedOp { fn verify(&self, ctx: &Context) -> Result<(), Error> { let op = &*self.get_operation().deref(ctx); + if let Some(alignment) = self + .get_attr_extern_alignment(ctx) + .map(|attr| attr.value().to_u64()) + && (alignment == 0 || !alignment.is_power_of_two()) + { + return verify_err!( + op.loc(), + "MirExternSharedOp alignment must be a non-zero power of two, found {}", + alignment + ); + } + // Check result is a shared memory pointer let res = op.get_result(0); let res_ty = res.get_type(ctx); @@ -1469,6 +1498,105 @@ impl Verify for MirExternSharedOp { } } +#[cfg(test)] +mod shared_alignment_verifier_tests { + use super::*; + use pliron::utils::apint::APInt; + use pliron::{ + builtin::{ + attributes::{IntegerAttr, TypeAttr}, + types::{IntegerType, Signedness}, + }, + common_traits::Verify, + operation::Operation, + }; + use std::num::NonZeroUsize; + + fn u64_attr(ctx: &Context, value: u64) -> IntegerAttr { + IntegerAttr::new( + IntegerType::get(ctx, 64, Signedness::Unsigned), + APInt::from_u64(value, NonZeroUsize::new(64).unwrap()), + ) + } + + fn shared_alloc(ctx: &mut Context, alignment: u64) -> MirSharedAllocOp { + let elem_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let result_ty = MirPtrType::get_shared(ctx, elem_ty, true); + let op = Operation::new( + ctx, + MirSharedAllocOp::get_concrete_op_info(), + vec![result_ty.into()], + vec![], + vec![], + 0, + ); + let alloc = MirSharedAllocOp::new(op); + alloc.set_attr_elem_type(ctx, TypeAttr::new(elem_ty)); + alloc.set_attr_size(ctx, u64_attr(ctx, 1)); + alloc.set_alignment_value(ctx, alignment); + alloc + } + + fn extern_shared(ctx: &mut Context, alignment: u64) -> MirExternSharedOp { + let elem_ty: TypeHandle = IntegerType::get(ctx, 32, Signedness::Unsigned).into(); + let result_ty = MirPtrType::get_shared(ctx, elem_ty, true); + let op = Operation::new( + ctx, + MirExternSharedOp::get_concrete_op_info(), + vec![result_ty.into()], + vec![], + vec![], + 0, + ); + let extern_shared = MirExternSharedOp::new(op); + extern_shared.set_alignment_value(ctx, alignment); + extern_shared + } + + /// The rejection must name the op and the offending value: that line is + /// what the user sees at the use site. + fn assert_rejected(result: Result<(), Error>, expected: &str) { + let error = result.expect_err("an invalid alignment must fail verification"); + let message = error.to_string(); + assert!( + message.contains(expected), + "expected `{expected}` in the diagnostic, got: {message}" + ); + } + + #[test] + fn shared_alloc_rejects_non_power_of_two_alignment() { + let mut ctx = Context::new(); + crate::register(&mut ctx); + + assert_rejected( + shared_alloc(&mut ctx, 3).verify(&ctx), + "MirSharedAllocOp alignment must be a non-zero power of two, found 3", + ); + assert_rejected( + shared_alloc(&mut ctx, 0).verify(&ctx), + "MirSharedAllocOp alignment must be a non-zero power of two, found 0", + ); + assert!(shared_alloc(&mut ctx, 8).verify(&ctx).is_ok()); + } + + #[test] + fn extern_shared_rejects_non_power_of_two_alignment() { + let mut ctx = Context::new(); + crate::register(&mut ctx); + + assert_rejected( + extern_shared(&mut ctx, 3).verify(&ctx), + "MirExternSharedOp alignment must be a non-zero power of two, found 3", + ); + assert_rejected( + extern_shared(&mut ctx, 0).verify(&ctx), + "MirExternSharedOp alignment must be a non-zero power of two, found 0", + ); + assert!(extern_shared(&mut ctx, 8).verify(&ctx).is_ok()); + } +} + /// Register memory operations into the given context. pub fn register(ctx: &mut Context) { MirAllocaOp::register(ctx); diff --git a/cuda-oxide/crates/dialect-mir/src/verification.rs b/cuda-oxide/crates/dialect-mir/src/verification.rs index 47c422a104..869d00ab4d 100644 --- a/cuda-oxide/crates/dialect-mir/src/verification.rs +++ b/cuda-oxide/crates/dialect-mir/src/verification.rs @@ -10,7 +10,6 @@ use pliron::{ builtin::ops::ConstantOp, context::{Context, Ptr}, linked_list::ContainsLinkedList, - location::Located, operation::Operation, result::Result, r#type::Typed, diff --git a/cuda-oxide/crates/dialect-mir/tests/ops_test/control_flow.rs b/cuda-oxide/crates/dialect-mir/tests/ops_test/control_flow.rs index cb8936f5fa..5384d397b5 100644 --- a/cuda-oxide/crates/dialect-mir/tests/ops_test/control_flow.rs +++ b/cuda-oxide/crates/dialect-mir/tests/ops_test/control_flow.rs @@ -7,8 +7,10 @@ use dialect_mir::ops::{MirAssertOp, MirCondBranchOp, MirFuncOp, MirGotoOp, MirRe use pliron::{ basic_block::BasicBlock, builtin::{ - attributes::TypeAttr, - op_interfaces::OperandSegmentInterface, + attributes::{OperandSegmentSizesAttr, StringAttr, TypeAttr}, + op_interfaces::{ + ATTR_KEY_OPERAND_SEGMENT_SIZES, BranchOpInterface, OperandSegmentInterface, + }, types::{FP32Type, FunctionType, IntegerType, Signedness}, }, common_traits::Verify, @@ -171,3 +173,248 @@ fn test_mir_control_flow_verify() { ); } } + +/// Conditional branches must reject malformed segment metadata before indexing it. +#[test] +fn test_mir_cond_branch_rejects_malformed_segments() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let source = BasicBlock::new(&mut ctx, None, vec![i1_ty.into(), i32_ty.into()]); + let condition = source.deref(&ctx).get_argument(0); + let value = source.deref(&ctx).get_argument(1); + let true_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let false_block = BasicBlock::new(&mut ctx, None, vec![]); + let branch = MirCondBranchOp::new(Operation::new( + &mut ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + vec![condition, value], + vec![true_block, false_block], + 0, + )); + + assert!(branch.verify(&ctx).is_err(), "Missing segment attribute"); + assert!(branch.verify_interfaces(&ctx).is_err()); + branch.get_operation().deref_mut(&ctx).attributes.set( + ATTR_KEY_OPERAND_SEGMENT_SIZES.clone(), + StringAttr::new("wrong attribute type".to_owned()), + ); + assert!(branch.verify(&ctx).is_err(), "Wrong segment attribute type"); + assert!(branch.verify_interfaces(&ctx).is_err()); + + for sizes in [ + vec![], + vec![2], + vec![1, 1], + vec![1, 1, 0, 0], + vec![1, 0, 0], + vec![0, 2, 0], + vec![2, 0, 0], + vec![1, 0, 1], + ] { + branch.set_operand_segment_sizes(&ctx, OperandSegmentSizesAttr(sizes.clone())); + assert!(branch.verify(&ctx).is_err(), "Accepted malformed {sizes:?}"); + assert!( + branch.verify_interfaces(&ctx).is_err(), + "Interfaces accepted {sizes:?}" + ); + } + + branch.set_operand_segment_sizes(&ctx, OperandSegmentSizesAttr(vec![1, 1, 0])); + assert!(branch.verify(&ctx).is_ok()); + assert!(branch.verify_interfaces(&ctx).is_ok()); +} + +/// Segment totals must not wrap to the operand count or panic during verification. +#[test] +fn test_mir_cond_branch_rejects_overflowing_segment_totals() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let source = BasicBlock::new(&mut ctx, None, vec![i1_ty.into(), i32_ty.into()]); + let condition = source.deref(&ctx).get_argument(0); + let value = source.deref(&ctx).get_argument(1); + let true_block = BasicBlock::new(&mut ctx, None, vec![i32_ty.into()]); + let false_block = BasicBlock::new(&mut ctx, None, vec![]); + let branch = MirCondBranchOp::new(Operation::new( + &mut ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + vec![condition, value], + vec![true_block, false_block], + 0, + )); + + // Both totals equal two modulo 2^32, matching the actual operand count. + for sizes in [vec![1, u32::MAX, 2], vec![u32::MAX, 1, 2]] { + branch.set_operand_segment_sizes(&ctx, OperandSegmentSizesAttr(sizes)); + assert!(branch.verify(&ctx).is_err()); + assert!(branch.verify_successor_operand_layout(&ctx).is_err()); + assert!(::verify(&branch, &ctx).is_err()); + assert!(branch.verify_interfaces(&ctx).is_err()); + } +} + +/// Invalid successor counts must be rejected without asking for an invalid range. +#[test] +fn test_mir_branches_reject_wrong_successor_counts() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let source = BasicBlock::new(&mut ctx, None, vec![i1_ty.into()]); + let condition = source.deref(&ctx).get_argument(0); + let target = BasicBlock::new(&mut ctx, None, vec![]); + + for count in [0, 2] { + let branch = MirGotoOp::new(Operation::new( + &mut ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + vec![], + vec![target; count], + 0, + )); + assert!(branch.verify(&ctx).is_err()); + assert!(branch.verify_interfaces(&ctx).is_err()); + } + for count in [0, 1, 3] { + let branch = MirCondBranchOp::new(Operation::new( + &mut ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + vec![condition], + vec![target; count], + 0, + )); + branch.set_operand_segment_sizes(&ctx, OperandSegmentSizesAttr(vec![1, 0, 0])); + assert!(branch.verify(&ctx).is_err()); + assert!(branch.verify_interfaces(&ctx).is_err()); + } +} + +/// Range-based mutation must preserve the condition and the other successor's operands. +#[test] +fn test_mir_cond_branch_successor_operand_mutation() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let i1_ty = IntegerType::get(&ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let f32_ty = FP32Type::get(&ctx); + let source = BasicBlock::new( + &mut ctx, + None, + vec![i1_ty.into(), i32_ty.into(), i32_ty.into(), f32_ty.into()], + ); + let condition = source.deref(&ctx).get_argument(0); + let first = source.deref(&ctx).get_argument(1); + let replacement = source.deref(&ctx).get_argument(2); + let second = source.deref(&ctx).get_argument(3); + let true_block = BasicBlock::new(&mut ctx, None, vec![]); + let false_block = BasicBlock::new(&mut ctx, None, vec![]); + let branch = MirCondBranchOp::new(Operation::new( + &mut ctx, + MirCondBranchOp::get_concrete_op_info(), + vec![], + vec![condition], + vec![true_block, false_block], + 0, + )); + branch.set_operand_segment_sizes(&ctx, OperandSegmentSizesAttr(vec![1, 0, 0])); + assert!(branch.verify(&ctx).is_ok()); + assert!(branch.verify_interfaces(&ctx).is_ok()); + assert_eq!(branch.successor_operand_range(&ctx, 0), 1..1); + assert_eq!(branch.successor_operand_range(&ctx, 1), 1..1); + + // Add false first: inserting true later must move, rather than overwrite, false. + assert_eq!(branch.add_successor_operand(&mut ctx, 1, second), 0); + BasicBlock::push_argument(false_block, &ctx, f32_ty.into()); + assert_eq!(branch.add_successor_operand(&mut ctx, 0, first), 0); + BasicBlock::push_argument(true_block, &ctx, i32_ty.into()); + assert_eq!(branch.successor_operand_range(&ctx, 0), 1..2); + assert_eq!(branch.successor_operand_range(&ctx, 1), 2..3); + assert_eq!(branch.successor_operands(&ctx, 0), vec![first]); + assert_eq!(branch.successor_operands(&ctx, 1), vec![second]); + assert!(branch.verify(&ctx).is_ok()); + assert!(branch.verify_interfaces(&ctx).is_ok()); + + branch.set_successor_operand(&ctx, 0, 0, second); + assert!(branch.verify(&ctx).is_err(), "Wrong true argument type"); + assert!(branch.verify_interfaces(&ctx).is_err()); + branch.set_successor_operand(&ctx, 0, 0, replacement); + branch.set_successor_operand(&ctx, 1, 0, first); + assert!(branch.verify(&ctx).is_err(), "Wrong false argument type"); + assert!(branch.verify_interfaces(&ctx).is_err()); + branch.set_successor_operand(&ctx, 1, 0, second); + assert_eq!(branch.successor_operands(&ctx, 0), vec![replacement]); + assert_eq!(branch.successor_operands(&ctx, 1), vec![second]); + assert_eq!(branch.get_operation().deref(&ctx).get_operand(0), condition); + + assert_eq!(branch.remove_successor_operand(&mut ctx, 0, 0), replacement); + BasicBlock::pop_argument(true_block, &ctx); + assert_eq!(branch.successor_operand_range(&ctx, 1), 1..2); + assert_eq!(branch.successor_operands(&ctx, 1), vec![second]); + assert!(branch.verify(&ctx).is_ok()); + assert!(branch.verify_interfaces(&ctx).is_ok()); + assert_eq!(branch.remove_successor_operand(&mut ctx, 1, 0), second); + BasicBlock::pop_argument(false_block, &ctx); + assert_eq!( + branch + .get_operation() + .deref(&ctx) + .operands() + .collect::>(), + vec![condition] + ); + assert!(branch.verify(&ctx).is_ok()); + assert!(branch.verify_interfaces(&ctx).is_ok()); +} + +/// Goto ranges must support empty and nonempty successors and type-checked mutation. +#[test] +fn test_mir_goto_successor_operand_mutation() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signed); + let f32_ty = FP32Type::get(&ctx); + let source = BasicBlock::new( + &mut ctx, + None, + vec![i32_ty.into(), i32_ty.into(), f32_ty.into()], + ); + let first = source.deref(&ctx).get_argument(0); + let replacement = source.deref(&ctx).get_argument(1); + let wrong_type = source.deref(&ctx).get_argument(2); + let target = BasicBlock::new(&mut ctx, None, vec![]); + let branch = MirGotoOp::new(Operation::new( + &mut ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + vec![], + vec![target], + 0, + )); + assert_eq!(branch.successor_operand_range(&ctx, 0), 0..0); + assert!(branch.verify(&ctx).is_ok()); + assert!(branch.verify_interfaces(&ctx).is_ok()); + assert_eq!(branch.add_successor_operand(&mut ctx, 0, first), 0); + assert!( + branch.verify(&ctx).is_err(), + "Successor argument is missing" + ); + BasicBlock::push_argument(target, &ctx, i32_ty.into()); + assert_eq!(branch.successor_operand_range(&ctx, 0), 0..1); + branch.set_successor_operand(&ctx, 0, 0, wrong_type); + assert!(branch.verify(&ctx).is_err()); + assert!(branch.verify_interfaces(&ctx).is_err()); + branch.set_successor_operand(&ctx, 0, 0, replacement); + assert_eq!(branch.successor_operands(&ctx, 0), vec![replacement]); + assert!(branch.verify(&ctx).is_ok()); + assert!(branch.verify_interfaces(&ctx).is_ok()); + assert_eq!(branch.remove_successor_operand(&mut ctx, 0, 0), replacement); + BasicBlock::pop_argument(target, &ctx); + assert!(branch.verify(&ctx).is_ok()); + assert!(branch.verify_interfaces(&ctx).is_ok()); +} diff --git a/cuda-oxide/crates/dialect-mir/tests/ops_test/scalars_and_funcs.rs b/cuda-oxide/crates/dialect-mir/tests/ops_test/scalars_and_funcs.rs index d39db3e74e..e45837dd4d 100644 --- a/cuda-oxide/crates/dialect-mir/tests/ops_test/scalars_and_funcs.rs +++ b/cuda-oxide/crates/dialect-mir/tests/ops_test/scalars_and_funcs.rs @@ -4,7 +4,7 @@ */ use dialect_mir::{ - attributes::{MirCastKindAttr, ReferenceParamValidityAttr}, + attributes::{InlineIntentAttr, MirCastKindAttr, ReferenceParamValidityAttr}, ops::{ MirAddOp, MirAssignOp, MirCallOp, MirCastOp, MirCheckedAddOp, MirCmpOp, MirConstantOp, MirDivOp, MirEqOp, MirFuncOp, MirGeOp, MirGtOp, MirLeOp, MirLtOp, MirMulOp, MirNeOp, @@ -385,6 +385,102 @@ fn test_mir_func_verify() { ); } +/// Function attributes must remain parseable without leaking their indentation +/// into the function body or later operations. +#[test] +fn function_attributes_print_and_parse_with_scoped_indentation() { + use pliron::{ + builtin::op_interfaces::SymbolOpInterface, + linked_list::ContainsLinkedList, + parsable::parse_from_str, + printable::{Printable, State}, + }; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let ty = FunctionType::get(&ctx, vec![], vec![]); + let op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let function = MirFuncOp::new(&mut ctx, op, TypeAttr::new(ty.into())); + function.set_symbol_name(&mut ctx, "roundtrip".try_into().unwrap()); + op.deref_mut(&ctx).attributes.set( + "note".try_into().unwrap(), + StringAttr::new("quoted \"text\"\nnext line\\tail".into()), + ); + let entry = BasicBlock::new(&mut ctx, None, vec![]); + entry.insert_at_back(op.deref(&ctx).get_region(0), &ctx); + let return_op = Operation::new( + &mut ctx, + dialect_mir::ops::MirReturnOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 0, + ); + return_op.insert_at_back(entry, &ctx); + op.verify(&ctx).unwrap(); + let state = State::default(); + state.set_indent_width(4); + let printed = function.print(&ctx, &state).to_string(); + assert_eq!(state.current_indent(), 0); + assert!(printed.contains("\n ["), "{printed}"); + + let mut parsed_ctx = Context::new(); + dialect_mir::register(&mut parsed_ctx); + let parsed = parse_from_str(Operation::top_level_parser(), &mut parsed_ctx, &printed) + .expect("printed function attributes must parse"); + assert_eq!(parsed.deref(&parsed_ctx).num_regions(), 1); + assert!( + parsed + .deref(&parsed_ctx) + .get_region(0) + .deref(&parsed_ctx) + .get_head() + .is_some() + ); + assert_eq!( + parsed + .deref(&parsed_ctx) + .attributes + .get::(&"note".try_into().unwrap()), + op.deref(&ctx) + .attributes + .get::(&"note".try_into().unwrap()), + ); + parsed + .verify(&parsed_ctx) + .expect("parsed function must verify"); +} + +/// Parsing a declaration supplies its empty body once; printing it without +/// extra attributes must still produce a valid function on the next parse. +#[test] +fn function_declaration_without_attributes_round_trips() { + use pliron::{parsable::parse_from_str, printable::Printable}; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let ty = FunctionType::get(&ctx, vec![], vec![]); + let text = format!("mir.func @declaration : {} []", ty.disp(&ctx)); + let parsed = parse_from_str(Operation::top_level_parser(), &mut ctx, &text).unwrap(); + assert_eq!(parsed.deref(&ctx).num_regions(), 1); + parsed.verify(&ctx).unwrap(); + let printed = parsed.disp(&ctx).to_string(); + + let mut parsed_ctx = Context::new(); + dialect_mir::register(&mut parsed_ctx); + let reparsed = parse_from_str(Operation::top_level_parser(), &mut parsed_ctx, &printed) + .expect("function without extra attributes must parse"); + assert_eq!(reparsed.deref(&parsed_ctx).num_regions(), 1); + reparsed.verify(&parsed_ctx).unwrap(); +} + #[test] fn test_mir_assign_verify() { let mut ctx = Context::new(); @@ -510,45 +606,123 @@ fn test_mir_func_reference_param_validity_verify() { }; let valid = make_func(&mut ctx, true); - valid.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr(4)); - valid.set_reference_param_validity(&mut ctx, 1, ReferenceParamValidityAttr(4)); + valid.set_reference_param_validity( + &mut ctx, + 0, + ReferenceParamValidityAttr { + dereferenceable_bytes: 4, + ..ReferenceParamValidityAttr::new(4) + }, + ); + valid.set_reference_param_validity(&mut ctx, 1, ReferenceParamValidityAttr::new(4)); assert!( valid.verify(&ctx).is_ok(), "Rust reference facts on kernel source arguments are valid" ); let non_kernel = make_func(&mut ctx, false); - non_kernel.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr(4)); + non_kernel.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr::new(4)); assert!( non_kernel.verify(&ctx).is_err(), "reference validity is kernel-entry-only" ); let raw = make_func(&mut ctx, true); - raw.set_reference_param_validity(&mut ctx, 2, ReferenceParamValidityAttr(4)); + raw.set_reference_param_validity(&mut ctx, 2, ReferenceParamValidityAttr::new(4)); assert!( raw.verify(&ctx).is_err(), "raw pointers cannot carry Rust-reference validity" ); let scalar = make_func(&mut ctx, true); - scalar.set_reference_param_validity(&mut ctx, 3, ReferenceParamValidityAttr(4)); + scalar.set_reference_param_validity(&mut ctx, 3, ReferenceParamValidityAttr::new(4)); assert!( scalar.verify(&ctx).is_err(), "non-pointer parameters cannot carry Rust-reference validity" ); let bad_alignment = make_func(&mut ctx, true); - bad_alignment.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr(3)); + bad_alignment.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr::new(3)); assert!( bad_alignment.verify(&ctx).is_err(), "alignment must be a non-zero power of two" ); let out_of_range = make_func(&mut ctx, true); - out_of_range.set_reference_param_validity(&mut ctx, 9, ReferenceParamValidityAttr(4)); + out_of_range.set_reference_param_validity(&mut ctx, 9, ReferenceParamValidityAttr::new(4)); assert!( out_of_range.verify(&ctx).is_err(), "source argument index must be in range" ); } + +/// An inline request is a callee property: every request is valid on an +/// ordinary function, none on a kernel entry, and a value of another attribute +/// type under the key is an error rather than "no request". +#[test] +fn test_mir_func_inline_intent_verify() { + use dialect_mir::ops::INLINE_INTENT_ATTR_KEY; + + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + + let make_func = |ctx: &mut Context, kernel: bool| { + let func_ty = FunctionType::get(ctx, vec![], vec![]); + let op = Operation::new( + ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func = MirFuncOp::new(ctx, op, TypeAttr::new(func_ty.into())); + if kernel { + func.get_operation().deref_mut(ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".to_string()), + ); + } + let region = func.get_operation().deref(ctx).get_region(0); + BasicBlock::new(ctx, None, vec![]).insert_at_front(region, ctx); + func + }; + + for intent in InlineIntentAttr::ALL { + let helper = make_func(&mut ctx, false); + helper.set_inline_intent(&mut ctx, intent); + assert!( + helper.verify(&ctx).is_ok(), + "{intent:?} is valid on a callee" + ); + assert_eq!(helper.inline_intent(&ctx), Some(intent)); + + let kernel = make_func(&mut ctx, true); + kernel.set_inline_intent(&mut ctx, intent); + assert!( + kernel.verify(&ctx).is_err(), + "{intent:?} must be rejected on a kernel entry" + ); + } + + let mistyped = make_func(&mut ctx, false); + mistyped.get_operation().deref_mut(&ctx).attributes.set( + INLINE_INTENT_ATTR_KEY.try_into().unwrap(), + StringAttr::new("noinline".to_string()), + ); + assert_eq!(mistyped.inline_intent(&ctx), None); + assert!( + mistyped.verify(&ctx).is_err(), + "a value of another attribute type is an error, never read as no request" + ); + + let plain = make_func(&mut ctx, false); + assert_eq!(plain.inline_intent(&ctx), None); + assert!(plain.verify(&ctx).is_ok()); + + let plain_kernel = make_func(&mut ctx, true); + assert!( + plain_kernel.verify(&ctx).is_ok(), + "a kernel without an intent is valid" + ); +} diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/asm.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/asm.rs index 6884178e96..2cbf0ac394 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/asm.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/asm.rs @@ -10,7 +10,6 @@ use pliron::{ builtin::attributes::{BoolAttr, StringAttr}, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/atomic.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/atomic.rs index a1d6bce7d4..ba6c34d54f 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/atomic.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/atomic.rs @@ -39,7 +39,6 @@ use pliron::{ common_traits::Verify, context::{Context, Ptr}, derive::{op_interface, op_interface_impl}, - location::Located, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/cluster.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/cluster.rs index dd4b5020e8..d4e6d93126 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/cluster.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/cluster.rs @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Result, diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/debug.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/debug.rs index a7c0f68fb4..a41b763c9f 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/debug.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/debug.rs @@ -26,7 +26,6 @@ use pliron::{ common_traits::Verify, context::Context, context::Ptr, - location::Located, op::Op, operation::Operation, result::Error, 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 449ea2c5e7..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 1c52191590..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Cluster Launch Control operations. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 fe423c1171..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -14,7 +14,6 @@ use pliron::{ builtin::op_interfaces::{NOpdsInterface, NResultsInterface}, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 488275d4d4..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 bfb3bed78c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for classic `cp.async` instructions. @@ -21,7 +21,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 d82e88c6aa..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated PTX debug controls. @@ -19,7 +19,6 @@ use pliron::{ common_traits::Verify, context::{Context, Ptr}, identifier::Identifier, - location::Located, op::Op, operation::Operation, result::Error, 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 ae2ab10cf8..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 666b1b6c3f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated warp leader-election operation. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 ec29e13a32..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -19,7 +19,6 @@ use pliron::{ common_traits::Verify, context::{Context, Ptr}, identifier::Identifier, - location::Located, op::Op, operation::Operation, result::Error, 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 693b61f9c3..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 93e4b44632..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 fac1180302..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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`. @@ -18,7 +18,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 38578d4237..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 ec97c54ac3..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -120,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( @@ -272,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( @@ -457,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( @@ -494,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 3c2463da32..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. mod active_mask; mod clc; 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 e7a3136723..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 30c46117b3..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 b7ac0ce129..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -18,7 +18,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 d992cb6310..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated packed conversion. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -432,6 +431,55 @@ impl Verify for CvtRnReluF16x2F32Op { } } +/// Converts packed bf16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order. +/// +/// The single input carries both packed lanes, and lane order is preserved. +#[pliron_op( + name = "nvvm.cvt_rn_satfinite_e4m3x2_bf16x2", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], +)] +pub struct CvtRnSatfiniteE4m3x2Bf16x2Op; + +impl CvtRnSatfiniteE4m3x2Bf16x2Op { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context, packed: Value) -> Ptr { + let result_ty = IntegerType::get(ctx, 16, Signedness::Unsigned); + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.into()], + vec![packed], + vec![], + 0, + ) + } +} + +impl Verify for CvtRnSatfiniteE4m3x2Bf16x2Op { + 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.cvt_rn_satfinite_e4m3x2_bf16x2 requires one operand and one result" + ); + } + if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32) + || !is_integer_width(ctx, op.get_result(0).get_type(ctx), 16) + { + return verify_err!( + op.loc(), + "nvvm.cvt_rn_satfinite_e4m3x2_bf16x2 requires one 32-bit integer operand and one 16-bit integer result" + ); + } + Ok(()) + } +} + /// Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation, preserving half order. /// /// The single input carries both packed lanes, and lane order is preserved. @@ -531,6 +579,55 @@ impl Verify for CvtRnSatfiniteE4m3x2F32Op { } } +/// Converts packed bf16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order. +/// +/// The single input carries both packed lanes, and lane order is preserved. +#[pliron_op( + name = "nvvm.cvt_rn_satfinite_e5m2x2_bf16x2", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], +)] +pub struct CvtRnSatfiniteE5m2x2Bf16x2Op; + +impl CvtRnSatfiniteE5m2x2Bf16x2Op { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context, packed: Value) -> Ptr { + let result_ty = IntegerType::get(ctx, 16, Signedness::Unsigned); + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.into()], + vec![packed], + vec![], + 0, + ) + } +} + +impl Verify for CvtRnSatfiniteE5m2x2Bf16x2Op { + 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.cvt_rn_satfinite_e5m2x2_bf16x2 requires one operand and one result" + ); + } + if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32) + || !is_integer_width(ctx, op.get_result(0).get_type(ctx), 16) + { + return verify_err!( + op.loc(), + "nvvm.cvt_rn_satfinite_e5m2x2_bf16x2 requires one 32-bit integer operand and one 16-bit integer result" + ); + } + Ok(()) + } +} + /// Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation, preserving half order. /// /// The single input carries both packed lanes, and lane order is preserved. @@ -630,6 +727,55 @@ impl Verify for CvtRnSatfiniteE5m2x2F32Op { } } +/// Converts packed bf16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order. +/// +/// The single input carries both packed lanes, and lane order is preserved. +#[pliron_op( + name = "nvvm.cvt_rn_satfinite_relu_e4m3x2_bf16x2", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], +)] +pub struct CvtRnSatfiniteReluE4m3x2Bf16x2Op; + +impl CvtRnSatfiniteReluE4m3x2Bf16x2Op { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context, packed: Value) -> Ptr { + let result_ty = IntegerType::get(ctx, 16, Signedness::Unsigned); + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.into()], + vec![packed], + vec![], + 0, + ) + } +} + +impl Verify for CvtRnSatfiniteReluE4m3x2Bf16x2Op { + 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.cvt_rn_satfinite_relu_e4m3x2_bf16x2 requires one operand and one result" + ); + } + if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32) + || !is_integer_width(ctx, op.get_result(0).get_type(ctx), 16) + { + return verify_err!( + op.loc(), + "nvvm.cvt_rn_satfinite_relu_e4m3x2_bf16x2 requires one 32-bit integer operand and one 16-bit integer result" + ); + } + Ok(()) + } +} + /// Converts packed f16x2 to packed e4m3x2 with nearest-even finite saturation and ReLU, preserving half order. /// /// The single input carries both packed lanes, and lane order is preserved. @@ -729,6 +875,55 @@ impl Verify for CvtRnSatfiniteReluE4m3x2F32Op { } } +/// Converts packed bf16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order. +/// +/// The single input carries both packed lanes, and lane order is preserved. +#[pliron_op( + name = "nvvm.cvt_rn_satfinite_relu_e5m2x2_bf16x2", + format, + interfaces = [NOpdsInterface<1>, NResultsInterface<1>], +)] +pub struct CvtRnSatfiniteReluE5m2x2Bf16x2Op; + +impl CvtRnSatfiniteReluE5m2x2Bf16x2Op { + pub fn new(op: Ptr) -> Self { + Self { op } + } + + pub fn build(ctx: &mut Context, packed: Value) -> Ptr { + let result_ty = IntegerType::get(ctx, 16, Signedness::Unsigned); + Operation::new( + ctx, + Self::get_concrete_op_info(), + vec![result_ty.into()], + vec![packed], + vec![], + 0, + ) + } +} + +impl Verify for CvtRnSatfiniteReluE5m2x2Bf16x2Op { + 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.cvt_rn_satfinite_relu_e5m2x2_bf16x2 requires one operand and one result" + ); + } + if !is_integer_width(ctx, op.get_operand(0).get_type(ctx), 32) + || !is_integer_width(ctx, op.get_result(0).get_type(ctx), 16) + { + return verify_err!( + op.loc(), + "nvvm.cvt_rn_satfinite_relu_e5m2x2_bf16x2 requires one 32-bit integer operand and one 16-bit integer result" + ); + } + Ok(()) + } +} + /// Converts packed f16x2 to packed e5m2x2 with nearest-even finite saturation and ReLU, preserving half order. /// /// The single input carries both packed lanes, and lane order is preserved. @@ -937,12 +1132,16 @@ pub(super) fn register(ctx: &mut Context) { CvtRnReluF16x2E4m3x2Op::register(ctx); CvtRnReluF16x2E5m2x2Op::register(ctx); CvtRnReluF16x2F32Op::register(ctx); + CvtRnSatfiniteE4m3x2Bf16x2Op::register(ctx); CvtRnSatfiniteE4m3x2F16x2Op::register(ctx); CvtRnSatfiniteE4m3x2F32Op::register(ctx); + CvtRnSatfiniteE5m2x2Bf16x2Op::register(ctx); CvtRnSatfiniteE5m2x2F16x2Op::register(ctx); CvtRnSatfiniteE5m2x2F32Op::register(ctx); + CvtRnSatfiniteReluE4m3x2Bf16x2Op::register(ctx); CvtRnSatfiniteReluE4m3x2F16x2Op::register(ctx); CvtRnSatfiniteReluE4m3x2F32Op::register(ctx); + CvtRnSatfiniteReluE5m2x2Bf16x2Op::register(ctx); CvtRnSatfiniteReluE5m2x2F16x2Op::register(ctx); CvtRnSatfiniteReluE5m2x2F32Op::register(ctx); CvtRzBf16x2F32Op::register(ctx); 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 ea9a88d3ac..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 2236ae8a42..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 552a0cb3c3..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -56,6 +55,7 @@ pub enum RegisterMmaKindAttr { Standard, F8f6f4, Mxf8f6f4, + Mxf4, } #[pliron_attr(name = "nvvm.register_mma_accumulator", format, verifier = "succ")] @@ -3731,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 726a2256e7..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated scalar arithmetic. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 1d2c361ae5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 424fd246dc..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! One structural operation for generated scalar math. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 578e3a3e4e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -19,7 +19,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 ed4f1c62d1..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -14,7 +14,6 @@ use pliron::{ builtin::types::IntegerType, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 c3e587573f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 02d2db9500..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 990baf92b6..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Tensor Core Generation 5 operations. @@ -19,7 +19,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 b6457435ee..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated Tensor Memory Accelerator operations. @@ -227,6 +227,21 @@ impl CpAsyncBulkTensorG2sTile1dOp { } } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +#[pliron_op( + name = "nvvm.cp_async_bulk_tensor_g2s_cta_tile_1d", + format, + verifier = "succ", + interfaces = [NOpdsInterface<5>, NResultsInterface<0>], +)] +pub struct CpAsyncBulkTensorG2sCtaTile1dOp; + +impl CpAsyncBulkTensorG2sCtaTile1dOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } +} + /// Starts a TMA tile copy from shared to global memory. #[pliron_op( name = "nvvm.cp_async_bulk_tensor_s2g_tile_1d", @@ -257,6 +272,21 @@ impl CpAsyncBulkTensorG2sTile2dOp { } } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +#[pliron_op( + name = "nvvm.cp_async_bulk_tensor_g2s_cta_tile_2d", + format, + verifier = "succ", + interfaces = [NOpdsInterface<6>, NResultsInterface<0>], +)] +pub struct CpAsyncBulkTensorG2sCtaTile2dOp; + +impl CpAsyncBulkTensorG2sCtaTile2dOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } +} + /// Starts a multicast TMA tile copy from global to cluster shared memory. #[pliron_op( name = "nvvm.cp_async_bulk_tensor_g2s_tile_2d_multicast", @@ -317,6 +347,21 @@ impl CpAsyncBulkTensorG2sTile3dOp { } } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +#[pliron_op( + name = "nvvm.cp_async_bulk_tensor_g2s_cta_tile_3d", + format, + verifier = "succ", + interfaces = [NOpdsInterface<7>, NResultsInterface<0>], +)] +pub struct CpAsyncBulkTensorG2sCtaTile3dOp; + +impl CpAsyncBulkTensorG2sCtaTile3dOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } +} + /// Starts a TMA tile copy from shared to global memory. #[pliron_op( name = "nvvm.cp_async_bulk_tensor_s2g_tile_3d", @@ -347,6 +392,21 @@ impl CpAsyncBulkTensorG2sTile4dOp { } } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +#[pliron_op( + name = "nvvm.cp_async_bulk_tensor_g2s_cta_tile_4d", + format, + verifier = "succ", + interfaces = [NOpdsInterface<8>, NResultsInterface<0>], +)] +pub struct CpAsyncBulkTensorG2sCtaTile4dOp; + +impl CpAsyncBulkTensorG2sCtaTile4dOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } +} + /// Starts a TMA tile copy from shared to global memory. #[pliron_op( name = "nvvm.cp_async_bulk_tensor_s2g_tile_4d", @@ -377,6 +437,21 @@ impl CpAsyncBulkTensorG2sTile5dOp { } } +/// Starts a TMA tile copy from global to the issuing CTA's shared memory. +#[pliron_op( + name = "nvvm.cp_async_bulk_tensor_g2s_cta_tile_5d", + format, + verifier = "succ", + interfaces = [NOpdsInterface<9>, NResultsInterface<0>], +)] +pub struct CpAsyncBulkTensorG2sCtaTile5dOp; + +impl CpAsyncBulkTensorG2sCtaTile5dOp { + pub fn new(op: Ptr) -> Self { + Self { op } + } +} + /// Starts a TMA tile copy from shared to global memory. #[pliron_op( name = "nvvm.cp_async_bulk_tensor_s2g_tile_5d", @@ -1697,16 +1772,21 @@ pub(super) fn register(ctx: &mut Context) { CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2Op::register(ctx); CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2CacheHintOp::register(ctx); CpAsyncBulkTensorG2sTile1dOp::register(ctx); + CpAsyncBulkTensorG2sCtaTile1dOp::register(ctx); CpAsyncBulkTensorS2gTile1dOp::register(ctx); CpAsyncBulkTensorG2sTile2dOp::register(ctx); + CpAsyncBulkTensorG2sCtaTile2dOp::register(ctx); CpAsyncBulkTensorG2sTile2dMulticastOp::register(ctx); CpAsyncBulkTensorG2sTile2dMulticastCg2Op::register(ctx); CpAsyncBulkTensorS2gTile2dOp::register(ctx); CpAsyncBulkTensorG2sTile3dOp::register(ctx); + CpAsyncBulkTensorG2sCtaTile3dOp::register(ctx); CpAsyncBulkTensorS2gTile3dOp::register(ctx); CpAsyncBulkTensorG2sTile4dOp::register(ctx); + CpAsyncBulkTensorG2sCtaTile4dOp::register(ctx); CpAsyncBulkTensorS2gTile4dOp::register(ctx); CpAsyncBulkTensorG2sTile5dOp::register(ctx); + CpAsyncBulkTensorG2sCtaTile5dOp::register(ctx); CpAsyncBulkTensorS2gTile5dOp::register(ctx); CpAsyncBulkTensorReduceAddIm2col3dOp::register(ctx); CpAsyncBulkTensorReduceAddIm2col4dOp::register(ctx); 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 def9f11d20..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 32aab1739b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operation for generated warp synchronization. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 c0c01f5ee1..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 767391efc9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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 9872b0e4b9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Structural operations for generated WGMMA controls. @@ -16,7 +16,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, diff --git a/cuda-oxide/crates/dialect-nvvm/src/ops/memory.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/memory.rs index b84e42f262..fd60ff734a 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/memory.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/memory.rs @@ -18,7 +18,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, @@ -75,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)) { @@ -91,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/src/ops/wgmma.rs b/cuda-oxide/crates/dialect-nvvm/src/ops/wgmma.rs index bd8f3faf7b..f5b02defd3 100644 --- a/cuda-oxide/crates/dialect-nvvm/src/ops/wgmma.rs +++ b/cuda-oxide/crates/dialect-nvvm/src/ops/wgmma.rs @@ -36,7 +36,6 @@ use pliron::{ }, common_traits::Verify, context::{Context, Ptr}, - location::Located, op::Op, operation::Operation, result::Error, 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/dialect-ptx/src/ops.rs b/cuda-oxide/crates/dialect-ptx/src/ops.rs index 807376d484..e9102f5b99 100644 --- a/cuda-oxide/crates/dialect-ptx/src/ops.rs +++ b/cuda-oxide/crates/dialect-ptx/src/ops.rs @@ -23,7 +23,6 @@ use pliron::{ common_traits::Verify, context::{Context, Ptr}, linked_list::ContainsLinkedList, - location::Located, op::Op, operation::Operation, region::Region, diff --git a/cuda-oxide/crates/llvm-export/README.md b/cuda-oxide/crates/llvm-export/README.md index 86ca7d2a81..98381cdc5f 100644 --- a/cuda-oxide/crates/llvm-export/README.md +++ b/cuda-oxide/crates/llvm-export/README.md @@ -5,8 +5,8 @@ pure-Rust textual LLVM IR (`.ll`) exporter, targeting the NVPTX backend. The LLVM dialect *modeling* (ops like `llvm.add`, types like `llvm.ptr`, attributes, and op-interfaces) lives upstream in the `pliron-llvm` crate. -cuda-oxide consumes it as a git dependency and no longer models the LLVM -dialect locally. This crate re-exports `pliron-llvm` so existing +cuda-oxide consumes it as a published crates.io release and no longer models +the LLVM dialect locally. This crate re-exports `pliron-llvm` so existing `llvm_export::{ops,types,attributes,op_interfaces}` paths keep resolving, adds a few GPU-specific extensions pliron-llvm lacks, and keeps the textual `.ll` exporter local (pliron-llvm only emits real `.ll` via an `llvm-sys` bridge, diff --git a/cuda-oxide/crates/llvm-export/src/export/debug.rs b/cuda-oxide/crates/llvm-export/src/export/debug.rs index 5b21a70906..ca0bfcfc31 100644 --- a/cuda-oxide/crates/llvm-export/src/export/debug.rs +++ b/cuda-oxide/crates/llvm-export/src/export/debug.rs @@ -1309,11 +1309,11 @@ mod tests { #[test] fn source_position_from_location_unwraps_named_locations() { - let mut ctx = Context::new(); + let ctx = Context::new(); let loc = Location::Named { name: "lowered".to_string(), child_loc: Box::new(Location::SrcPos { - src: Source::new_from_file(&mut ctx, PathBuf::from("/tmp/kernel.rs")), + src: Source::new_from_file(&ctx, PathBuf::from("/tmp/kernel.rs")), pos: SourcePosition { line: 12, column: 4, diff --git a/cuda-oxide/crates/llvm-export/src/export/function.rs b/cuda-oxide/crates/llvm-export/src/export/function.rs index 324afd7996..db8900f9a2 100644 --- a/cuda-oxide/crates/llvm-export/src/export/function.rs +++ b/cuda-oxide/crates/llvm-export/src/export/function.rs @@ -16,14 +16,13 @@ use pliron::{ attribute::Attribute, basic_block::BasicBlock, builtin::{ - attributes::{FPDoubleAttr, FPSingleAttr, IntegerAttr, TypeAttr}, - op_interfaces::{BranchOpInterface, SymbolOpInterface}, + attributes::{FPDoubleAttr, FPSingleAttr, IntegerAttr, OperandSegmentSizesAttr, TypeAttr}, + op_interfaces::{ATTR_KEY_OPERAND_SEGMENT_SIZES, BranchOpInterface, SymbolOpInterface}, type_interfaces::FunctionTypeInterface, types::IntegerType, }, - context::Ptr, + context::{Context, Ptr}, linked_list::ContainsLinkedList, - location::Located, op::Op, operation::Operation, printable::Printable, @@ -61,6 +60,103 @@ fn sanitize_comment(text: &str) -> String { .collect() } +/// Check conditional-branch layout before Pliron's infallible range accessors. +/// +/// Public export callers need not have verified the source module. Pliron 0.19 +/// also sums segment sizes as u32 without checking overflow, so validate the +/// complete layout here before it can index or allocate from malformed counts. +/// Looks for a pointer into shared (address space 3) or local (address space +/// 5) memory inside a kernel parameter type, through struct fields and array +/// elements. On a hit it returns the address space and appends the path that +/// reaches the pointer to `path`: `.N` for struct field `N` and `[0]` for an +/// array element (every element has the same type, so the first stands for +/// all of them). A bare pointer leaves `path` empty. +fn device_only_pointer_in(ctx: &Context, ty: TypeHandle, path: &mut String) -> Option { + let ty_ref = ty.deref(ctx); + if let Some(pointer) = ty_ref.downcast_ref::() { + let space = pointer.address_space(); + return matches!(space, 3 | 5).then_some(space); + } + if let Some(array) = ty_ref.downcast_ref::() { + let mark = path.len(); + path.push_str("[0]"); + if let Some(space) = device_only_pointer_in(ctx, array.elem_type(), path) { + return Some(space); + } + path.truncate(mark); + return None; + } + if let Some(structure) = ty_ref.downcast_ref::() { + if structure.is_opaque() { + return None; + } + for (field_index, field) in structure.fields().enumerate() { + let mark = path.len(); + path.push_str(&format!(".{field_index}")); + if let Some(space) = device_only_pointer_in(ctx, field, path) { + return Some(space); + } + path.truncate(mark); + } + } + None +} + +fn validate_conditional_branch_layout(ctx: &Context, branch: &ops::CondBrOp) -> Result<(), String> { + let operation = branch.get_operation().deref(ctx); + let sizes = operation + .attributes + .get::(&ATTR_KEY_OPERAND_SEGMENT_SIZES) + .ok_or_else(|| { + "conditional branch requires an operand-segment-sizes attribute".to_owned() + })?; + if !matches!(sizes.0.as_slice(), [1, _, _]) { + return Err( + "conditional branch requires three operand segments with exactly one condition" + .to_owned(), + ); + } + let total = sizes + .0 + .iter() + .try_fold(0_u32, |total, size| total.checked_add(*size)) + .ok_or_else(|| "conditional branch operand segment total overflows u32".to_owned())?; + if total as usize != operation.get_num_operands() { + return Err(format!( + "conditional branch operand segments describe {total} values, but it has {} operands", + operation.get_num_operands() + )); + } + if operation.get_num_successors() != 2 { + return Err("conditional branch requires exactly two successors".to_owned()); + } + Ok(()) +} + +/// Spell only the facts already proven at the import boundary. Shared by +/// definitions and declarations, including the empty-region declaration path. +fn export_reference_parameter_attrs( + validity: ops::KernelReferenceParamValidityAttr, + include_alignment: bool, + output: &mut String, +) { + write!(output, " nonnull").unwrap(); + if validity.pointee_alignment > 1 && include_alignment { + write!(output, " align {}", validity.pointee_alignment).unwrap(); + } + if validity.no_alias { + write!(output, " noalias").unwrap(); + } + if validity.dereferenceable_bytes > 0 { + write!( + output, + " dereferenceable({})", + validity.dereferenceable_bytes + ) + .unwrap(); + } +} + /// Final LLVM spelling of a pliron function symbol. /// /// Shared-global owner attributes carry the raw symbol, so the module prepass, @@ -200,9 +296,19 @@ impl<'a> ModuleExportState<'a> { } else { "global" }; + // Host-written storage (`#[constant]`): the initializer is only the + // load-time image. Without this, whole-program LTO (libNVVM + // `-gen-lto` + nvJitLink `-lto`) internalizes the never-stored + // global, folds its loads to the initializer and drops the symbol + // the host looks up with `cuModuleGetGlobal`. + let externally_initialized = if global.is_host_written(self.ctx) { + "externally_initialized " + } else { + "" + }; write!( output, - "@{name} = addrspace({address_space}) {storage_keyword} " + "@{name} = addrspace({address_space}) {externally_initialized}{storage_keyword} " ) .unwrap(); self.export_type(ty, output)?; @@ -799,7 +905,7 @@ impl<'a> ModuleExportState<'a> { "kernel `@{fixed_func_name}` reference validity parameter {index} is not an LLVM pointer" )); } - reference_param_validities[index] = Some(validity.0); + reference_param_validities[index] = Some(validity); } let is_declaration = func.get_operation().deref(self.ctx).regions().count() == 0; @@ -930,23 +1036,40 @@ impl<'a> ModuleExportState<'a> { // parameter type today; if a frontend type ever surfaces it at a // kernel boundary, it needs a new arm here for the same reason as // plain shared. + // + // The check looks through struct fields and array elements as well. + // A struct or array parameter reaches the kernel as the bytes the + // host wrote, so a shared or local pointer inside one is as + // impossible to supply as a bare one. Before this recursion a + // `struct { p: SharedPtr, n: u32 }` parameter compiled and the + // kernel applied `cvta.shared.u64` to host bytes. if is_kernel { for (index, arg_ty) in func_ty.arg_types().iter().enumerate() { - let arg_ref = arg_ty.deref(self.ctx); - let Some(pointer) = arg_ref.downcast_ref::() else { + let mut path = String::new(); + let Some(address_space) = device_only_pointer_in(self.ctx, *arg_ty, &mut path) + else { continue; }; - let space = match pointer.address_space() { - 3 => "shared", - 5 => "local", - _ => continue, + let space = if address_space == 3 { + "shared" + } else { + "local" }; + let reason = format!( + "{space} memory is allocated by the device at launch, so the host has no \ + address to pass, and the driver refuses a module whose entry declares one" + ); + if path.is_empty() { + return Err(format!( + "kernel `@{fixed_func_name}` parameter {index} is a pointer into {space} \ + memory (address space {address_space}); {reason}" + )); + } + // The field path keeps this message apart from the bare form, + // so a fixture can pin the nested case on its own. return Err(format!( - "kernel `@{fixed_func_name}` parameter {index} is a pointer into {space} \ - memory (address space {}); {space} memory is allocated by the device at \ - launch, so the host has no address to pass, and the driver refuses a \ - module whose entry declares one", - pointer.address_space() + "kernel `@{fixed_func_name}` parameter {index} holds a pointer into {space} \ + memory (address space {address_space}) at field path `{path}`; {reason}" )); } } @@ -975,11 +1098,8 @@ impl<'a> ModuleExportState<'a> { let grid_constant = grid_constant_params .iter() .find(|parameter| parameter.index == i); - if let Some(alignment) = reference_param_validities[i] { - write!(output, " nonnull").unwrap(); - if alignment > 1 && grid_constant.is_none() { - write!(output, " align {alignment}").unwrap(); - } + if let Some(validity) = reference_param_validities[i] { + export_reference_parameter_attrs(validity, grid_constant.is_none(), output); } if let Some(parameter) = grid_constant { self.export_grid_constant_parameter_attrs( @@ -1015,16 +1135,8 @@ impl<'a> ModuleExportState<'a> { .iter(self.ctx) .next(); - // Check for alwaysinline attribute (from #[inline(always)]). - // Emitted as a function attribute keyword between the parameter - // list and the body open brace. - let alwaysinline_key: pliron::identifier::Identifier = "alwaysinline".try_into().unwrap(); - let is_alwaysinline = attrs - .get::(&alwaysinline_key) - .is_some(); - if let Some(entry_block) = entry_block_opt { - let func_loc = func.get_operation().deref(self.ctx).loc(); + let func_loc = func.get_operation().deref(self.ctx).loc().clone(); let debug_name = ops::debug_function_name(self.ctx, func.get_operation()) .unwrap_or_else(|| fixed_func_name.clone()); let debug_scope = @@ -1046,15 +1158,9 @@ impl<'a> ModuleExportState<'a> { let block = entry_block.deref(self.ctx); let args = block.arguments(); let mut entry_prologue = String::new(); - // Kernel Rust references may carry importer-proven `nonnull` and - // pointee `align N`. Every other parameter remains bare. - // - // In particular this deliberately does NOT infer `noalias`, - // `readonly`, `dereferenceable`, or any validity fact from LLVM - // pointer shape. `DisjointSlice` and raw pointers therefore retain - // the conservative behavior required by their unsafe construction - // contracts. Future alias attributes need their own audited proof - // source rather than extending this mechanical export step. + // Kernel Rust references carry only importer-proven facts. + // LLVM pointer shape alone grants no reference promises, so raw + // pointers, slice metadata, and other parameters remain bare. for (i, arg) in args.enumerate() { if i > 0 { write!(output, ", ").unwrap(); @@ -1064,11 +1170,8 @@ impl<'a> ModuleExportState<'a> { let grid_constant = grid_constant_params .iter() .find(|parameter| parameter.index == i); - if let Some(alignment) = reference_param_validities[i] { - write!(output, " nonnull").unwrap(); - if alignment > 1 && grid_constant.is_none() { - write!(output, " align {alignment}").unwrap(); - } + if let Some(validity) = reference_param_validities[i] { + export_reference_parameter_attrs(validity, grid_constant.is_none(), output); } if let Some(parameter) = grid_constant { self.export_grid_constant_parameter_attrs( @@ -1113,12 +1216,9 @@ impl<'a> ModuleExportState<'a> { // gets its `bar.sync.aligned` pushed into a `tid`-dependent branch // and deadlocks. opt's FunctionAttrs strips `convergent` from // functions it proves never reach a convergent op. - // alwaysinline (from #[inline(always)]) and !dbg are independent: - // either, both, or neither can be present. Emit the inline keyword - // before the convergent attr group #0, then the debug scope. - if is_alwaysinline { - write!(output, " alwaysinline").unwrap(); - } + // Function attributes (an `#[inline]` keyword among them, set by + // `mir-lower` in `llvm_func_attrs`) precede the convergent attr + // group #0, then the debug scope. if let Some(scope_id) = debug_scope { writeln!(output, " #0 !dbg !{scope_id} {{").unwrap(); } else { @@ -1337,6 +1437,7 @@ impl<'a> ModuleExportState<'a> { validate_edge(block, dest, &args)?; pred_map.entry(dest).or_default().push((block, args)); } else if let Some(branch) = term_dyn.downcast_ref::() { + validate_conditional_branch_layout(self.ctx, branch)?; let succs: Vec<_> = term.deref(self.ctx).successors().collect(); let true_dest = succs[0]; let false_dest = succs[1]; @@ -1409,11 +1510,8 @@ impl<'a> ModuleExportState<'a> { let grid_constant = grid_constant_params .iter() .find(|parameter| parameter.index == i); - if let Some(alignment) = reference_param_validities[i] { - write!(output, " nonnull").unwrap(); - if alignment > 1 && grid_constant.is_none() { - write!(output, " align {alignment}").unwrap(); - } + if let Some(validity) = reference_param_validities[i] { + export_reference_parameter_attrs(validity, grid_constant.is_none(), output); } if let Some(parameter) = grid_constant { self.export_grid_constant_parameter_attrs( diff --git a/cuda-oxide/crates/llvm-export/src/export/module.rs b/cuda-oxide/crates/llvm-export/src/export/module.rs index a99893e363..1a9e934851 100644 --- a/cuda-oxide/crates/llvm-export/src/export/module.rs +++ b/cuda-oxide/crates/llvm-export/src/export/module.rs @@ -16,7 +16,6 @@ use pliron::{ }, context::Context, linked_list::ContainsLinkedList, - location::Located, op::Op, operation::Operation, r#type::{TypeHandle, Typed}, @@ -208,7 +207,7 @@ fn prepare_debug_shared_function_scopes( let raw_name = function.get_symbol_name(state.ctx); let owner = super::function::exported_function_name(raw_name.as_ref()); if state.debug_shared_function_scopes.contains_key(&owner) { - let loc = function.get_operation().deref(state.ctx).loc(); + let loc = function.get_operation().deref(state.ctx).loc().clone(); let debug_name = ops::debug_function_name(state.ctx, function.get_operation()) .unwrap_or_else(|| owner.clone()); state.debug_subprogram_for_function(&debug_name, &owner, &loc); diff --git a/cuda-oxide/crates/llvm-export/src/export/ops.rs b/cuda-oxide/crates/llvm-export/src/export/ops.rs index 42a5be6f96..1c1216bc43 100644 --- a/cuda-oxide/crates/llvm-export/src/export/ops.rs +++ b/cuda-oxide/crates/llvm-export/src/export/ops.rs @@ -19,7 +19,6 @@ use pliron::{ types::{FP32Type, FP64Type, IntegerType}, }, context::Ptr, - location::Located, op::Op, operation::Operation, value::Value, @@ -486,10 +485,10 @@ impl<'a> ModuleExportState<'a> { Some(LlvmOp::Constant(op)) => self.emit_constant(op, value_names)?, Some(LlvmOp::AddressOf(op)) => self.emit_address_of(op, value_names, output)?, Some(LlvmOp::DebugValue(op)) => { - self.emit_debug_value(op, value_names, debug_scope, &op_loc, output)? + self.emit_debug_value(op, value_names, debug_scope, op_loc, output)? } Some(LlvmOp::DebugValueList(op)) => { - self.emit_debug_value_list(op, value_names, debug_scope, &op_loc, output)? + self.emit_debug_value_list(op, value_names, debug_scope, op_loc, output)? } // Unknown None if self.legacy_typed_pointers() => { @@ -511,12 +510,12 @@ impl<'a> ModuleExportState<'a> { output, output_before, debug_scope, - &op_loc, + op_loc, allow_scope_debug_fallback, ); } if let Some(alloca) = debug_alloca { - self.emit_debug_declare_for_alloca(alloca, value_names, debug_scope, &op_loc, output)?; + self.emit_debug_declare_for_alloca(alloca, value_names, debug_scope, op_loc, output)?; } Ok(()) diff --git a/cuda-oxide/crates/llvm-export/src/lib.rs b/cuda-oxide/crates/llvm-export/src/lib.rs index 6509d9277c..ba5dc7d272 100644 --- a/cuda-oxide/crates/llvm-export/src/lib.rs +++ b/cuda-oxide/crates/llvm-export/src/lib.rs @@ -210,16 +210,31 @@ pub mod ops { /// Kernel-entry parameter validity proven at the Rust MIR import boundary. /// - /// Presence proves `nonnull`; the payload is the rustc ABI alignment of - /// the pointee represented by this physical LLVM parameter. It deliberately - /// carries no aliasing, readonly, or dereferenceability promise. + /// Presence proves `nonnull`; each field transports an independently + /// proven promise for this physical LLVM parameter. A zero byte count + /// records no dereferenceability promise. #[pliron_attr( name = "llvm.kernel_reference_param_validity", - format = "$0", + format, verifier = "succ" )] #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] - pub struct KernelReferenceParamValidityAttr(pub u64); + pub struct KernelReferenceParamValidityAttr { + pub pointee_alignment: u64, + pub no_alias: bool, + pub dereferenceable_bytes: u64, + } + + impl KernelReferenceParamValidityAttr { + /// Non-nullness and alignment, without additional reference promises. + pub const fn new(pointee_alignment: u64) -> Self { + Self { + pointee_alignment, + no_alias: false, + dereferenceable_bytes: 0, + } + } + } fn kernel_reference_param_validity_key(index: usize) -> Identifier { Identifier::try_new(reserved_oxide_symbols::kernel_reference_param_validity_key( @@ -271,10 +286,10 @@ pub mod ops { "kernel reference parameter validity attribute `{key_text}` has the wrong attribute type" )); }; - if validity.0 == 0 || !validity.0.is_power_of_two() { + if validity.pointee_alignment == 0 || !validity.pointee_alignment.is_power_of_two() { return Err(format!( "kernel reference parameter {index} alignment must be a non-zero power of two, found {}", - validity.0 + validity.pointee_alignment )); } result.push((index, validity)); @@ -411,8 +426,14 @@ pub mod ops { const GLOBAL_SHARED_SOURCE_NAME_KEY: &str = "cuda_oxide_global_shared_source_name"; /// Marks a `GlobalOp` as externally consumed even when device code does not /// reference it. The exporter adds marked globals to `@llvm.used`, keeping - /// profiler metadata alive through libNVVM and nvJitLink materialization. + /// profiler metadata and host-written constant memory alive through libNVVM + /// and nvJitLink materialization. const GLOBAL_RETAINED_KEY: &str = "cuda_oxide_global_retained"; + /// Marks a `GlobalOp` whose contents the host writes after module load + /// (`cuModuleGetGlobal` + copy). The exporter renders it + /// `externally_initialized`, so no optimizer, including libNVVM/nvJitLink + /// LTO, may treat the initializer as the value device code reads. + const GLOBAL_HOST_WRITTEN_KEY: &str = "cuda_oxide_global_host_written"; /// One pointer-width relocation inside an evaluated Rust static initializer. /// @@ -2549,6 +2570,10 @@ pub mod ops { fn mark_retained(&self, ctx: &mut Context); /// Whether this global was explicitly marked as externally consumed. fn is_retained(&self, ctx: &Context) -> bool; + /// Declare that the host writes this global after module load. + fn mark_host_written(&self, ctx: &mut Context); + /// Whether the host writes this global after module load. + fn is_host_written(&self, ctx: &Context) -> bool; } impl GlobalOpExt for GlobalOp { @@ -2670,6 +2695,25 @@ pub mod ops { .get::(&key) .is_some() } + + fn mark_host_written(&self, ctx: &mut Context) { + let key = + Identifier::try_new(GLOBAL_HOST_WRITTEN_KEY.to_string()).expect("valid identifier"); + self.get_operation() + .deref_mut(ctx) + .attributes + .set(key, pliron::builtin::attributes::UnitAttr); + } + + fn is_host_written(&self, ctx: &Context) -> bool { + let key = + Identifier::try_new(GLOBAL_HOST_WRITTEN_KEY.to_string()).expect("valid identifier"); + self.get_operation() + .deref(ctx) + .attributes + .get::(&key) + .is_some() + } } #[cfg(test)] diff --git a/cuda-oxide/crates/llvm-export/tests/export_test/control_flow.rs b/cuda-oxide/crates/llvm-export/tests/export_test/control_flow.rs index 33f9fc3cc3..4a94078212 100644 --- a/cuda-oxide/crates/llvm-export/tests/export_test/control_flow.rs +++ b/cuda-oxide/crates/llvm-export/tests/export_test/control_flow.rs @@ -11,7 +11,8 @@ use llvm_export::{ use pliron::{ basic_block::BasicBlock, builtin::{ - attributes::IntegerAttr, + attributes::{IntegerAttr, OperandSegmentSizesAttr, UnitAttr}, + op_interfaces::{ATTR_KEY_OPERAND_SEGMENT_SIZES, OperandSegmentInterface}, ops::ModuleOp, types::{IntegerType, Signedness}, }, @@ -24,6 +25,125 @@ use std::num::NonZero; use crate::common::module_top_block; +/// Build a conditional branch forwarding one i32 argument to each destination. +fn conditional_branch_module(ctx: &mut Context) -> (ModuleOp, CondBrOp) { + let module = ModuleOp::new(ctx, "conditional_segments".try_into().unwrap()); + let module_block = module_top_block(ctx, &module); + let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let func_ty = FuncType::get( + ctx, + i32_ty.into(), + vec![i1_ty.into(), i32_ty.into(), i32_ty.into()], + false, + ); + let func = FuncOp::new(ctx, "conditional".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(ctx); + let condition = entry.deref(ctx).get_argument(0); + let if_true = entry.deref(ctx).get_argument(1); + let if_false = entry.deref(ctx).get_argument(2); + let region = func.get_operation().deref(ctx).get_region(0); + let true_dest = BasicBlock::new(ctx, None, vec![i32_ty.into()]); + let false_dest = BasicBlock::new(ctx, None, vec![i32_ty.into()]); + for destination in [true_dest, false_dest] { + destination.insert_at_back(region, ctx); + let result = destination.deref(ctx).get_argument(0); + ReturnOp::new(ctx, Some(result)) + .get_operation() + .insert_at_back(destination, ctx); + } + let branch = CondBrOp::new( + ctx, + condition, + true_dest, + vec![if_true], + false_dest, + vec![if_false], + ); + branch.get_operation().insert_at_back(entry, ctx); + func.get_operation().insert_at_back(module_block, ctx); + (module, branch) +} + +/// Export through the public entry point without assuming prior source verification. +fn export_conditional_module(ctx: &Context, module: &ModuleOp) -> Result { + export_module_to_string_with_config( + ctx, + module, + &NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7), + ) +} + +/// Invalid segment totals must return an error before accessing successor operands. +#[test] +fn exporter_rejects_out_of_bounds_conditional_segments() { + let mut ctx = Context::new(); + let (module, branch) = conditional_branch_module(&mut ctx); + branch.set_operand_segment_sizes(&ctx, OperandSegmentSizesAttr(vec![1, 3, 0])); + assert!(export_conditional_module(&ctx, &module).is_err()); +} + +/// A missing segment attribute must not reach Pliron's infallible segment accessor. +#[test] +fn exporter_rejects_missing_conditional_segments() { + let mut ctx = Context::new(); + let (module, branch) = conditional_branch_module(&mut ctx); + branch + .get_operation() + .deref_mut(&ctx) + .attributes + .0 + .remove(&ATTR_KEY_OPERAND_SEGMENT_SIZES); + assert!(export_conditional_module(&ctx, &module).is_err()); +} + +/// The segment attribute must have its expected typed representation. +#[test] +fn exporter_rejects_wrong_conditional_segment_attribute_type() { + let mut ctx = Context::new(); + let (module, branch) = conditional_branch_module(&mut ctx); + branch + .get_operation() + .deref_mut(&ctx) + .attributes + .set(ATTR_KEY_OPERAND_SEGMENT_SIZES.clone(), UnitAttr); + assert!(export_conditional_module(&ctx, &module).is_err()); +} + +/// Oversized counts must fail without overflowing either total or prefix arithmetic. +#[test] +fn exporter_rejects_overflowing_conditional_segments() { + let mut ctx = Context::new(); + let (module, branch) = conditional_branch_module(&mut ctx); + branch.set_operand_segment_sizes(&ctx, OperandSegmentSizesAttr(vec![1, u32::MAX, 2])); + assert!(export_conditional_module(&ctx, &module).is_err()); +} + +/// Conditional branches need three segments, starting with exactly one condition. +#[test] +fn exporter_rejects_invalid_conditional_segment_shapes() { + for sizes in [vec![], vec![1, 2], vec![0, 2, 1], vec![2, 0, 1]] { + let mut ctx = Context::new(); + let (module, branch) = conditional_branch_module(&mut ctx); + branch.set_operand_segment_sizes(&ctx, OperandSegmentSizesAttr(sizes.clone())); + assert!( + export_conditional_module(&ctx, &module).is_err(), + "{sizes:?}" + ); + } +} + +/// Valid segment metadata keeps both conditional edges and their distinct forwarded values. +#[test] +fn exporter_preserves_valid_conditional_segments() { + let mut ctx = Context::new(); + let (module, _) = conditional_branch_module(&mut ctx); + let ir = export_conditional_module(&ctx, &module).unwrap(); + assert!(ir.contains("br i1 %v0, label %bb0, label %bb1"), "{ir}"); + assert!(ir.contains("phi i32 [ %v1, %entry ]"), "{ir}"); + assert!(ir.contains("phi i32 [ %v2, %entry ]"), "{ir}"); +} + #[test] fn exporter_rejects_extra_predecessor_values_before_emitting_phis() { let mut ctx = Context::new(); 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 d338f5eb60..abc11b7af6 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 @@ -30,7 +30,7 @@ use pliron::{ context::Context, identifier::Identifier, linked_list::ContainsLinkedList, - location::{Located, Location}, + location::Location, op::Op, utils::apint::APInt, }; @@ -830,6 +830,8 @@ fn full_debug_metadata_describes_as4_constant_with_semantic_rust_type() { 4, ); constant.set_address_space(&mut ctx, llvm_export::types::address_space::CONSTANT); + // `create_device_global` marks every `#[constant]` as host-written. + constant.mark_host_written(&mut ctx); constant.set_initializer_hex(&mut ctx, "00000000"); llvm_export::ops::set_debug_global_variable( &mut ctx, @@ -918,7 +920,8 @@ fn full_debug_metadata_describes_as4_constant_with_semantic_rust_type() { .find(|line| line.starts_with(&format!("@{linkage_name} = "))) .expect("constant-memory definition"); assert!( - definition.contains("addrspace(4) global [4 x i8] c\"\\00\\00\\00\\00\"") + definition + .contains("addrspace(4) externally_initialized global [4 x i8] c\"\\00\\00\\00\\00\"") && definition.contains(", align 4, !dbg !"), "physical AS4 byte storage must retain a global debug attachment:\n{full}" ); diff --git a/cuda-oxide/crates/llvm-export/tests/export_test/grid_constant.rs b/cuda-oxide/crates/llvm-export/tests/export_test/grid_constant.rs index ce6404520e..e6628f0993 100644 --- a/cuda-oxide/crates/llvm-export/tests/export_test/grid_constant.rs +++ b/cuda-oxide/crates/llvm-export/tests/export_test/grid_constant.rs @@ -23,6 +23,7 @@ use pliron::{ }, common_traits::Verify, context::Context, + identifier::Identifier, linked_list::ContainsLinkedList, op::Op, operation::Operation, @@ -496,12 +497,10 @@ fn grid_constant_rejects_invalid_parameter_contracts() { ); } "incomplete" => { - attrs.0.remove( - &format!("{LLVM_GRID_CONSTANT_ALIGN_ATTR_PREFIX}0") - .as_str() - .try_into() - .unwrap(), - ); + let key = + Identifier::try_new(format!("{LLVM_GRID_CONSTANT_ALIGN_ATTR_PREFIX}0")) + .unwrap(); + attrs.0.remove(&key); } "wrong_attribute_kind" => { attrs.set( diff --git a/cuda-oxide/crates/llvm-export/tests/export_test/kernels_and_roots.rs b/cuda-oxide/crates/llvm-export/tests/export_test/kernels_and_roots.rs index e6490122cf..d80832d0a2 100644 --- a/cuda-oxide/crates/llvm-export/tests/export_test/kernels_and_roots.rs +++ b/cuda-oxide/crates/llvm-export/tests/export_test/kernels_and_roots.rs @@ -10,7 +10,7 @@ use llvm_export::{ export_module_with_externs_and_roots, }, ops::{FuncOp, GlobalOp, GlobalOpExt, ReturnOp}, - types::{FuncType, PointerType, VoidType}, + types::{ArrayType, FuncType, PointerType, StructLayout, StructType, VoidType}, }; use pliron::{ basic_block::BasicBlock, @@ -22,6 +22,7 @@ use pliron::{ context::{Context, Ptr}, identifier::Identifier, op::Op, + r#type::TypeHandle, }; use crate::common::module_top_block; @@ -109,6 +110,57 @@ fn legacy_llvm_used_roots_retained_address_space_global() { ); } +/// Host-written storage must not look like a known constant to LTO in any IR +/// dialect, and it stays rooted in `@llvm.used` so the symbol the host +/// resolves survives even when every read of it is optimized away. Ordinary +/// globals keep their plain definition and stay out of the root list. +#[test] +fn host_written_global_exports_externally_initialized() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "host_written".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + + // Mirror `create_device_global`: a `#[constant]` carries both markers. + let params = GlobalOp::new(&mut ctx, "PARAMS".try_into().unwrap(), i32_ty.into()); + params.set_address_space(&mut ctx, 4); + params.mark_retained(&mut ctx); + params.mark_host_written(&mut ctx); + params.get_operation().insert_at_back(module_block, &ctx); + + let ordinary = GlobalOp::new(&mut ctx, "ORDINARY".try_into().unwrap(), i32_ty.into()); + ordinary.set_address_space(&mut ctx, 1); + ordinary.get_operation().insert_at_back(module_block, &ctx); + + let opaque_used = "@llvm.used = appending global [1 x ptr] [ptr addrspacecast (ptr addrspace(4) @PARAMS to ptr)], section \"llvm.metadata\""; + let legacy_used = "@llvm.used = appending global [1 x i8*] [i8* addrspacecast (i32 addrspace(4)* @PARAMS to i8*)], section \"llvm.metadata\""; + let modern = NvvmExportConfig::new(NvvmIrDialect::Modern); + let legacy = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); + for (ir, used) in [ + ( + export_module_to_string_with_config(&ctx, &module, &PtxExportConfig), + opaque_used, + ), + ( + export_module_to_string_with_config(&ctx, &module, &modern), + opaque_used, + ), + ( + export_module_to_string_with_config(&ctx, &module, &legacy), + legacy_used, + ), + ] { + let ir = ir.expect("host-written global export succeeds"); + assert!( + ir.contains("@PARAMS = addrspace(4) externally_initialized global i32"), + "{ir}" + ); + assert!(ir.contains(used), "{ir}"); + assert!(ir.contains("@ORDINARY = addrspace(1) global i32"), "{ir}"); + assert!(!ir.contains("@ORDINARY to"), "{ir}"); + } +} + #[test] fn ptx_export_records_kernel_roots_for_internalization() { let mut ctx = Context::new(); @@ -195,10 +247,22 @@ fn pointer_param_function( address_spaces: &[u32], is_kernel: bool, ) { - let params: Vec<_> = address_spaces + let params: Vec = address_spaces .iter() .map(|space| PointerType::get(ctx, *space).into()) .collect(); + typed_param_function(ctx, module_block, name, params, is_kernel); +} + +/// Builds a `void` function with the given parameter types and an empty +/// body, optionally carrying the `gpu_kernel` marker. +fn typed_param_function( + ctx: &mut Context, + module_block: Ptr, + name: &str, + params: Vec, + is_kernel: bool, +) { let void_ty = VoidType::get(ctx); let func_ty = FuncType::get(ctx, void_ty.into(), params, false); let func = FuncOp::new(ctx, name.try_into().unwrap(), func_ty); @@ -256,6 +320,165 @@ fn kernel_rejects_a_local_memory_pointer_parameter() { ); } +/// The host fills a struct parameter byte by byte, so a shared pointer +/// inside one is as impossible to supply as a bare one. The exporter looks +/// through the aggregate and names the field path it found, with a message +/// distinct from the bare form. +#[test] +fn kernel_rejects_a_shared_memory_pointer_inside_a_struct_parameter() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "kernel_shared_struct_param".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let shared_ptr_ty = PointerType::get(&ctx, 3); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let args_ty = StructType::get_unnamed( + &ctx, + ( + vec![shared_ptr_ty.into(), i32_ty.into()], + StructLayout::Unpacked, + ), + ); + typed_param_function( + &mut ctx, + module_block, + "shared_struct_param", + vec![args_ty.into()], + true, + ); + + let error = export_module_to_string(&ctx, &module) + .expect_err("export must refuse a kernel parameter struct that holds a shared pointer"); + assert!( + error.contains( + "kernel `@shared_struct_param` parameter 0 holds a pointer into shared memory \ + (address space 3) at field path `.0`" + ), + "{error}" + ); +} + +/// An array parameter is refused the same way, and the path names the +/// element position. +#[test] +fn kernel_rejects_a_shared_memory_pointer_inside_an_array_parameter() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "kernel_shared_array_param".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let generic_ptr_ty = PointerType::get(&ctx, 0); + let shared_ptr_ty = PointerType::get(&ctx, 3); + let pair_ty = ArrayType::get(&ctx, shared_ptr_ty.into(), 2); + typed_param_function( + &mut ctx, + module_block, + "shared_array_param", + vec![generic_ptr_ty.into(), pair_ty.into()], + true, + ); + + let error = export_module_to_string(&ctx, &module) + .expect_err("export must refuse a kernel parameter array of shared pointers"); + assert!( + error.contains( + "kernel `@shared_array_param` parameter 1 holds a pointer into shared memory \ + (address space 3) at field path `[0]`" + ), + "{error}" + ); +} + +/// A pointer buried under a struct, an array and another struct is found and +/// reported with its full path; local memory is refused on the same grounds. +#[test] +fn kernel_rejects_a_local_memory_pointer_nested_below_an_array_element() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "kernel_local_nested_param".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let local_ptr_ty = PointerType::get(&ctx, 5); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let leaf_ty = StructType::get_unnamed( + &ctx, + ( + vec![local_ptr_ty.into(), i32_ty.into()], + StructLayout::Unpacked, + ), + ); + let pair_ty = ArrayType::get(&ctx, leaf_ty.into(), 2); + let args_ty = StructType::get_unnamed( + &ctx, + (vec![i32_ty.into(), pair_ty.into()], StructLayout::Unpacked), + ); + typed_param_function( + &mut ctx, + module_block, + "local_nested_param", + vec![args_ty.into()], + true, + ); + + let error = export_module_to_string(&ctx, &module) + .expect_err("export must refuse a nested local-memory pointer in a kernel parameter"); + assert!( + error.contains( + "kernel `@local_nested_param` parameter 0 holds a pointer into local memory \ + (address space 5) at field path `.1[0].0`" + ), + "{error}" + ); +} + +/// The recursion refuses only what the host cannot supply: a struct of +/// generic and global pointers stays a valid kernel parameter, and a device +/// (non-kernel) function may take the struct that holds a shared pointer. +#[test] +fn kernel_keeps_a_struct_of_host_addressable_pointers_and_device_functions_keep_nested_shared() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "kernel_allowed_struct_params".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let generic_ptr_ty = PointerType::get(&ctx, 0); + let global_ptr_ty = PointerType::get(&ctx, 1); + let shared_ptr_ty = PointerType::get(&ctx, 3); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + let host_ty = StructType::get_unnamed( + &ctx, + ( + vec![generic_ptr_ty.into(), global_ptr_ty.into()], + StructLayout::Unpacked, + ), + ); + let shared_ty = StructType::get_unnamed( + &ctx, + ( + vec![shared_ptr_ty.into(), i32_ty.into()], + StructLayout::Unpacked, + ), + ); + typed_param_function( + &mut ctx, + module_block, + "host_struct_param", + vec![host_ty.into()], + true, + ); + typed_param_function( + &mut ctx, + module_block, + "device_shared_struct", + vec![shared_ty.into()], + false, + ); + + let ir = export_module_to_string(&ctx, &module) + .expect("a host-addressable struct parameter and a device function's shared struct export"); + assert!( + ir.contains("define ptx_kernel void @host_struct_param({ ptr, ptr addrspace(1) } %v0)"), + "{ir}" + ); + assert!( + ir.contains("@device_shared_struct({ ptr addrspace(3), i32 } %v0)"), + "{ir}" + ); +} + /// The refusal is scoped to what the host cannot supply. Generic, global, /// and constant pointer parameters are host-addressable and stay accepted on /// kernels, and a device (non-kernel) function may carry any state space on diff --git a/cuda-oxide/crates/llvm-export/tests/export_test/metadata.rs b/cuda-oxide/crates/llvm-export/tests/export_test/metadata.rs index 8f8314d917..6656491855 100644 --- a/cuda-oxide/crates/llvm-export/tests/export_test/metadata.rs +++ b/cuda-oxide/crates/llvm-export/tests/export_test/metadata.rs @@ -20,7 +20,6 @@ use pliron::{ context::Context, identifier::Identifier, linked_list::ContainsLinkedList, - location::Located, op::Op, utils::apint::APInt, }; @@ -142,41 +141,48 @@ fn nvvm_metadata_version_uses_next_allocated_metadata_id() { } #[test] -fn export_alwaysinline_function_attribute_uses_llvm_define_syntax() { - let mut ctx = Context::new(); - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_block = module_top_block(&mut ctx, &module); - - let void_ty = VoidType::get(&ctx); - let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); - let func = FuncOp::new(&mut ctx, "inline_helper".try_into().unwrap(), func_ty); - let entry = func.get_or_create_entry_block(&mut ctx); - ReturnOp::new(&mut ctx, None) - .get_operation() - .insert_at_back(entry, &ctx); +fn export_inline_intent_function_attribute_uses_llvm_define_syntax() { + // `mir-lower` places an `#[inline]` keyword in `llvm_func_attrs`; the generic + // attribute renderer writes it after the parameter list, before attr group #0. + // The keywords are listed here as plain text: this crate renders whatever + // attribute it is handed and does not depend on the MIR dialect's enum. + for keyword in ["inlinehint", "alwaysinline", "noinline"] { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); - let key: pliron::identifier::Identifier = "alwaysinline".try_into().unwrap(); - func.get_operation() - .deref_mut(&ctx) - .attributes - .set(key, StringAttr::new("true".to_string())); - func.get_operation().insert_at_back(module_block, &ctx); + let void_ty = VoidType::get(&ctx); + let func_ty = FuncType::get(&ctx, void_ty.to_handle(), vec![], false); + let func = FuncOp::new(&mut ctx, "inline_helper".try_into().unwrap(), func_ty); + let entry = func.get_or_create_entry_block(&mut ctx); + ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + set_inline_keyword(&mut ctx, &func, keyword); + func.get_operation().insert_at_back(module_block, &ctx); - let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); - let define_line = ir - .lines() - .find(|line| line.starts_with("define void @inline_helper(")) - .expect("inline helper definition"); - assert_eq!( - define_line, "define void @inline_helper() alwaysinline #0 {", - "`alwaysinline` must be emitted after the parameter list, before attr group #0:\n{ir}" - ); - assert!( - ir.contains("attributes #0 = { convergent }"), - "convergent attribute group must still be emitted:\n{ir}" - ); + let ir = export_module_to_string(&ctx, &module).expect("export succeeds"); + let define_line = ir + .lines() + .find(|line| line.starts_with("define void @inline_helper(")) + .expect("inline helper definition"); + assert_eq!( + define_line, + format!("define void @inline_helper() {keyword} #0 {{"), + "{keyword} must be emitted after the parameter list, before attr group #0:\n{ir}" + ); + assert!( + ir.contains("attributes #0 = { convergent }"), + "convergent attribute group must still be emitted:\n{ir}" + ); + } } +fn set_inline_keyword(ctx: &mut Context, func: &FuncOp, keyword: &str) { + let mut attrs = llvm_export::ops::LlvmAttributesAttr::new(); + attrs.set(keyword, llvm_export::ops::LlvmAttrValue::Unit); + func.set_attr_llvm_func_attrs(ctx, attrs); +} #[test] fn export_alwaysinline_coexists_with_debug_scope() { // alwaysinline and the !dbg scope are emitted on the same define line and @@ -197,11 +203,8 @@ fn export_alwaysinline_coexists_with_debug_scope() { ret.get_operation().deref_mut(&ctx).set_loc(ret_loc); ret.get_operation().insert_at_back(entry, &ctx); - let key: pliron::identifier::Identifier = "alwaysinline".try_into().unwrap(); - func.get_operation() - .deref_mut(&ctx) - .attributes - .set(key, StringAttr::new("true".to_string())); + set_inline_keyword(&mut ctx, &func, "alwaysinline"); + func.get_operation().insert_at_back(module_block, &ctx); let config = DebugConfig { diff --git a/cuda-oxide/crates/llvm-export/tests/export_test/reference_validity.rs b/cuda-oxide/crates/llvm-export/tests/export_test/reference_validity.rs index 58b3a36924..4767ff0c28 100644 --- a/cuda-oxide/crates/llvm-export/tests/export_test/reference_validity.rs +++ b/cuda-oxide/crates/llvm-export/tests/export_test/reference_validity.rs @@ -62,13 +62,13 @@ fn kernel_reference_validity_exports_nonnull_and_alignment() { &mut ctx, func.get_operation(), 0, - KernelReferenceParamValidityAttr(4), + KernelReferenceParamValidityAttr::new(4), ); set_kernel_reference_param_validity( &mut ctx, func.get_operation(), 1, - KernelReferenceParamValidityAttr(1), + KernelReferenceParamValidityAttr::new(1), ); func.get_operation().insert_at_back(module_block, &ctx); @@ -93,7 +93,7 @@ fn legacy_kernel_reference_validity_follows_typed_pointer_spelling() { &mut ctx, func.get_operation(), 0, - KernelReferenceParamValidityAttr(4), + KernelReferenceParamValidityAttr::new(4), ); func.get_operation().insert_at_back(module_block, &ctx); @@ -116,7 +116,7 @@ fn reference_validity_rejects_non_kernel_non_pointer_and_out_of_range() { &mut ctx, non_kernel.get_operation(), 0, - KernelReferenceParamValidityAttr(4), + KernelReferenceParamValidityAttr::new(4), ); non_kernel .get_operation() @@ -132,7 +132,7 @@ fn reference_validity_rejects_non_kernel_non_pointer_and_out_of_range() { &mut ctx, scalar.get_operation(), 0, - KernelReferenceParamValidityAttr(4), + KernelReferenceParamValidityAttr::new(4), ); scalar.get_operation().insert_at_back(scalar_block, &ctx); let error = export_module_to_string_with_config(&ctx, &scalar_module, &PtxExportConfig) @@ -146,7 +146,7 @@ fn reference_validity_rejects_non_kernel_non_pointer_and_out_of_range() { &mut ctx, range.get_operation(), 3, - KernelReferenceParamValidityAttr(4), + KernelReferenceParamValidityAttr::new(4), ); range.get_operation().insert_at_back(range_block, &ctx); let error = export_module_to_string_with_config(&ctx, &range_module, &PtxExportConfig) @@ -165,13 +165,98 @@ fn reference_validity_rejects_malformed_marker_indices() { let ptr = PointerType::get(&ctx, 0); let func = kernel_with_params(&mut ctx, "malformed_marker", vec![ptr.into()], true); let key = format!("{KERNEL_REFERENCE_PARAM_VALIDITY_KEY_PREFIX}{suffix}"); - func.get_operation() - .deref_mut(&ctx) - .attributes - .set(key.try_into().unwrap(), KernelReferenceParamValidityAttr(4)); + func.get_operation().deref_mut(&ctx).attributes.set( + key.try_into().unwrap(), + KernelReferenceParamValidityAttr::new(4), + ); func.get_operation().insert_at_back(module_block, &ctx); let error = export_module_to_string_with_config(&ctx, &module, &PtxExportConfig) .expect_err("malformed reference-validity marker must fail"); assert!(error.contains("invalid parameter index"), "{error}"); } } + +#[test] +fn reference_validity_exports_each_fact_independently_in_all_dialects() { + use llvm_export::export::ExportBackendConfig; + + let modern_nvvm = NvvmExportConfig::new(NvvmIrDialect::Modern); + let legacy_nvvm = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); + let configs: [&dyn ExportBackendConfig; 3] = [&PtxExportConfig, &modern_nvvm, &legacy_nvvm]; + for (config, pointer) in configs.into_iter().zip(["ptr", "ptr", "i8*"]) { + // Cover definitions, declarations, and declarations with an empty region. + for shape in 0..3 { + let definition = shape == 0; + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "independent_facts".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let ptr = PointerType::get(&ctx, 0); + let integer = IntegerType::get(&ctx, 64, Signedness::Signless); + let mut params = vec![ptr.into(); 6]; + params.push(integer.into()); + let func = if definition { + kernel_with_params(&mut ctx, "independent_facts", params, true) + } else { + let ty = FuncType::get(&ctx, VoidType::get(&ctx).into(), params, false); + let func = FuncOp::new(&mut ctx, "independent_facts".try_into().unwrap(), ty); + if shape == 2 { + pliron::operation::Operation::add_region(func.get_operation(), &mut ctx); + } + func.get_operation().deref_mut(&ctx).attributes.set( + "gpu_kernel".try_into().unwrap(), + StringAttr::new("true".to_string()), + ); + func + }; + for (index, (alignment, no_alias, bytes)) in [ + (4, true, 4), + (8, true, 8), + (1, false, 0), + (16, false, 128), + (4, true, 0), + ] + .into_iter() + .enumerate() + { + set_kernel_reference_param_validity( + &mut ctx, + func.get_operation(), + index, + KernelReferenceParamValidityAttr { + pointee_alignment: alignment, + no_alias, + dereferenceable_bytes: bytes, + }, + ); + } + func.get_operation().insert_at_back(module_block, &ctx); + let ir = export_module_to_string_with_config(&ctx, &module, config).unwrap(); + assert!(!ir.contains(" readonly"), "{ir}"); + let expected = [ + format!("{pointer} nonnull align 4 noalias dereferenceable(4)"), + format!("{pointer} nonnull align 8 noalias dereferenceable(8)"), + format!("{pointer} nonnull"), + format!("{pointer} nonnull align 16 dereferenceable(128)"), + format!("{pointer} nonnull align 4 noalias"), + pointer.to_string(), + "i64".to_string(), + ] + .into_iter() + .enumerate() + .map(|(index, parameter)| { + if definition { + format!("{parameter} %v{index}") + } else { + parameter + } + }) + .collect::>() + .join(", "); + assert!( + ir.contains(&format!("@independent_facts({expected})")), + "{ir}" + ); + assert!(!ir.contains("dereferenceable(0)"), "{ir}"); + } + } +} 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 5d88f4a234..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, @@ -91,19 +91,17 @@ pub struct CollectedFunction { /// sidecar. Full debug uses supported events to keep affected pointer values /// in stack slots; other modes emit no extra code. pub statement_debug_info: Option, - /// True if the function is marked `#[inline(always)]` in rustc's - /// `CodegenFnAttrs`. The stable_mir API does not expose inline hints, so - /// this is queried via `rustc_middle::TyCtxt::codegen_fn_attrs` in - /// `rustc-codegen-cuda` and threaded through. + /// The function's `#[inline]` request from rustc's `CodegenFnAttrs`. The + /// stable_mir API does not expose inline attributes, so this is queried via + /// `rustc_middle::TyCtxt::codegen_fn_attrs` in `rustc-codegen-cuda` and + /// threaded through. /// - /// When true, the LLVM `alwaysinline` attribute is emitted on the - /// function definition. The existing matched LLVM middle-end (`opt -O2`), - /// when available, can then honor the attribute before PTX generation; - /// this flag does not add a separate mandatory inliner pass. - /// - /// This preserves Rust's inline intent for device helpers and avoids - /// making helper boundaries depend entirely on later optimizer heuristics. - pub is_inline_always: bool, + /// `mir-lower` turns the request into the matching LLVM function attribute + /// on the definition, so the matched LLVM middle-end (`opt -O2`) honors it + /// before PTX generation; no separate inliner pass is added. This keeps + /// helper boundaries, `#[inline(never)]` ones included, from depending + /// entirely on optimizer heuristics. + pub inline: Option, } /// Per-function statement debug records aligned with the stable MIR body. @@ -191,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). @@ -265,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, @@ -494,11 +501,16 @@ pub fn run_pipeline( config: &PipelineConfig, known_defs: crate::translator::facts::KnownDefs, ) -> Result { - // Install the driver-resolved lang-item ids for this run. Set (not - // merged) every entry: the ids are only valid inside the caller's - // `rustc_internal::run` context, so a stale set from a previous run on - // this thread must never survive. + // Install the driver-resolved lang-item ids and emitted-function identities + // for this run. Replace (rather than merge) both thread-local sets so + // stale state from a previous run on this thread cannot survive. crate::translator::facts::set_known_defs(known_defs); + crate::translator::facts::set_emitted_functions( + functions + .iter() + .map(|func| func.instance.mangled_name().to_string()) + .collect(), + ); prepare_output_dir(&config.output_dir)?; @@ -518,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 { @@ -545,7 +558,7 @@ pub fn run_pipeline( func.rustc_mir_block_count, &func.rustc_mono_successors, func.is_kernel, - func.is_inline_always, + func.inline, Some(&func.export_name), &mut legaliser, config.debug_kind, @@ -580,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); @@ -589,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); @@ -598,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 { @@ -642,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, + }), } } @@ -766,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); @@ -1296,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, @@ -1350,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, @@ -1442,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 1a701eccac..2af4b2c883 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/body.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/body.rs @@ -26,7 +26,7 @@ use crate::error::{TranslationErr, TranslationResult}; use crate::pipeline::StatementDebugInfoMap; use crate::translator::location::span_to_location; use crate::translator::values::{self, SlotAddrSpaceMap, ValueMap}; -use dialect_mir::attributes::ReferenceParamValidityAttr; +use dialect_mir::attributes::{InlineIntentAttr, ReferenceParamValidityAttr}; use dialect_mir::ops::MirFuncOp; use dialect_mir::types::address_space; use llvm_export::export::DebugKind; @@ -39,7 +39,6 @@ use pliron::basic_block::BasicBlock; use pliron::builtin::op_interfaces::SymbolOpInterface; use pliron::context::{Context, Ptr}; use pliron::identifier::{Identifier, Legaliser}; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::{input_err_noloc, input_error_noloc}; @@ -1881,8 +1880,8 @@ fn emit_entry_allocas( /// * `rustc_mono_successors` - Exact per-block successor edges computed by /// rustc's monomorphization rules under the device runtime-check policy /// * `is_kernel` - Add `gpu_kernel` attribute for kernel entry points -/// * `is_inline_always` - Add `alwaysinline` attribute (non-kernel functions -/// marked `#[inline(always)]` in rustc) +/// * `inline` - The `#[inline]` request, recorded as [`InlineIntentAttr`] on +/// non-kernel functions (`set_inline_intent_attr`) /// * `override_name` - Custom export name (defaults to instance name) pub fn translate_body( ctx: &mut Context, @@ -1891,7 +1890,7 @@ pub fn translate_body( rustc_mir_block_count: usize, rustc_mono_successors: &[Vec], is_kernel: bool, - is_inline_always: bool, + inline: Option, override_name: Option<&str>, legaliser: &mut Legaliser, debug_kind: DebugKind, @@ -2141,16 +2140,34 @@ pub fn translate_body( matches!(signature.skip_binder().abi, rustc_public::ty::Abi::Rust) }); if is_kernel && has_rust_abi { + let fn_abi = instance.fn_abi().map_err(|error| { + input_error_noloc!(TranslationErr::unsupported(format!( + "could not query kernel reference argument ABI: {error:?}" + ))) + })?; + if fn_abi.fixed_count as usize != num_args { + return input_err_noloc!(TranslationErr::invalid_op(format!( + "kernel has {num_args} source arguments but rustc FnAbi declares {}", + fn_abi.fixed_count + ))); + } 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) else { + 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; }; mir_func_op.set_reference_param_validity( ctx, arg_idx, - ReferenceParamValidityAttr(validity.pointee_alignment), + ReferenceParamValidityAttr { + pointee_alignment: validity.pointee_alignment, + no_alias: validity.no_alias, + dereferenceable_bytes: validity.dereferenceable_bytes, + }, ); } } @@ -2364,7 +2381,7 @@ pub fn translate_body( llvm_export::ops::set_debug_source_scope_map(ctx, op_ptr, scope_map); } - set_alwaysinline_attr_from_flag(ctx, &mir_func_op, is_kernel, is_inline_always); + set_inline_intent_attr(ctx, &mir_func_op, is_kernel, inline); // Get the function body region (region 0) let region_ptr = op_ptr.deref(ctx).get_region(0); @@ -2498,24 +2515,20 @@ fn function_debug_name(instance: &mono::Instance, is_kernel: bool, export_name: } } -/// Propagate `#[inline(always)]` as an LLVM `alwaysinline` function -/// attribute. Kernel entry points are excluded because they're `.entry` in PTX -/// and never callees, so marking them `alwaysinline` would be a no-op at best -/// and rejected by LLVM at worst. -fn set_alwaysinline_attr_from_flag( +/// Record the `#[inline]` request on the MIR function (`MirFuncOp::set_inline_intent`; +/// `mir-lower` makes it the LLVM function attribute). Kernel entry points are +/// excluded because they're `.entry` in PTX and never callees, so there is no +/// call the request could apply to. +fn set_inline_intent_attr( ctx: &mut Context, mir_func_op: &MirFuncOp, is_kernel: bool, - is_inline_always: bool, + inline: Option, ) { - if is_inline_always && !is_kernel { - let attr = pliron::builtin::attributes::StringAttr::new("true".to_string()); - let key: Identifier = "alwaysinline".try_into().unwrap(); - mir_func_op - .get_operation() - .deref_mut(ctx) - .attributes - .set(key, attr); + if let Some(intent) = inline + && !is_kernel + { + mir_func_op.set_inline_intent(ctx, intent); } } @@ -2731,78 +2744,83 @@ pub fn cuda_oxide_device_generated_kernel(mut wrapped: Wrapper) -> u32 { } #[test] - fn inline_always_flag_reaches_llvm_func_attr_before_export() { - let mut ctx = Context::new(); - crate::translator::register_dialects(&mut ctx); - - let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); - let module_op = module.get_operation(); - let module_region = module_op.deref(&ctx).get_region(0); - let module_block = { - let existing = { - let region = module_region.deref(&ctx); - region.iter(&ctx).next() + fn inline_intent_reaches_llvm_func_attr_before_export() { + for (intent, keyword) in [ + (InlineIntentAttr::Hint, "inlinehint"), + (InlineIntentAttr::Always, "alwaysinline"), + (InlineIntentAttr::Never, "noinline"), + ] { + let mut ctx = Context::new(); + crate::translator::register_dialects(&mut ctx); + + let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap()); + let module_op = module.get_operation(); + let module_region = module_op.deref(&ctx).get_region(0); + let module_block = { + let existing = { + let region = module_region.deref(&ctx); + region.iter(&ctx).next() + }; + if let Some(block) = existing { + block + } else { + let block = BasicBlock::new(&mut ctx, None, vec![]); + block.insert_at_back(module_region, &ctx); + block + } }; - if let Some(block) = existing { - block - } else { - let block = BasicBlock::new(&mut ctx, None, vec![]); - block.insert_at_back(module_region, &ctx); - block - } - }; - let func_type = FunctionType::get(&ctx, vec![], vec![]); - let func_type_attr = TypeAttr::new(func_type.into()); - let mir_func = { - let op = Operation::new( + let func_type = FunctionType::get(&ctx, vec![], vec![]); + let func_type_attr = TypeAttr::new(func_type.into()); + let mir_func = { + let op = Operation::new( + &mut ctx, + MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let func = MirFuncOp::new(&mut ctx, op, func_type_attr); + func.set_symbol_name(&mut ctx, "inline_helper".try_into().unwrap()); + func + }; + + set_inline_intent_attr(&mut ctx, &mir_func, false, Some(intent)); + llvm_export::ops::set_debug_function_name( &mut ctx, - MirFuncOp::get_concrete_op_info(), - vec![], - vec![], - vec![], - 1, + mir_func.get_operation(), + "source_crate::inline_helper", ); - let func = MirFuncOp::new(&mut ctx, op, func_type_attr); - func.set_symbol_name(&mut ctx, "inline_helper".try_into().unwrap()); - func - }; + mir_func.get_operation().insert_at_back(module_block, &ctx); - set_alwaysinline_attr_from_flag(&mut ctx, &mir_func, false, true); - llvm_export::ops::set_debug_function_name( - &mut ctx, - mir_func.get_operation(), - "source_crate::inline_helper", - ); - mir_func.get_operation().insert_at_back(module_block, &ctx); - - mir_lower::register(&mut ctx); - mir_lower::lower_mir_to_llvm(&mut ctx, module_op).expect("lowering succeeds"); - - let llvm_func = { - let block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); - block - .deref(&ctx) - .iter(&ctx) - .find_map(|op| Operation::get_op::(op, &ctx)) - .expect("lowered LLVM function") - }; + mir_lower::register(&mut ctx); + mir_lower::lower_mir_to_llvm(&mut ctx, module_op).expect("lowering succeeds"); - let key: Identifier = "alwaysinline".try_into().unwrap(); - assert!( - llvm_func - .get_operation() - .deref(&ctx) - .attributes - .0 - .contains_key(&key), - "`is_inline_always` must become an LLVM dialect alwaysinline attribute before export", - ); - assert_eq!( - llvm_export::ops::debug_function_name(&ctx, llvm_func.get_operation()).as_deref(), - Some("source_crate::inline_helper"), - "MIR-to-LLVM lowering must preserve the source-facing function name", - ); + let llvm_func = { + let block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + block + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("lowered LLVM function") + }; + + let keywords: Vec = llvm_func + .get_attr_llvm_func_attrs(&ctx) + .map(|attrs| attrs.iter().map(|(name, _)| name.to_string()).collect()) + .unwrap_or_default(); + assert_eq!( + keywords, + [keyword], + "{intent:?} must become its one LLVM function attribute before export", + ); + assert_eq!( + llvm_export::ops::debug_function_name(&ctx, llvm_func.get_operation()).as_deref(), + Some("source_crate::inline_helper"), + "MIR-to-LLVM lowering must preserve the source-facing function name", + ); + } } /// Exercise `debug_fragment` against composite `VarDebugInfo` produced by 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 2be0ceafe8..a6b24ac367 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/facts.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/facts.rs @@ -85,6 +85,12 @@ thread_local! { index_mut_trait: None, }) }; + // Exact monomorphized Rust functions that this pipeline run will emit. + // + // Store rustc's linkage identity rather than the legalised module symbol: + // call classification must not depend on Legaliser ordering or spelling. + static EMITTED_FUNCTIONS: std::cell::RefCell> = + const { std::cell::RefCell::new(std::collections::BTreeSet::new()) }; } /// Installs the driver-resolved lang-item ids for this pipeline run, @@ -99,6 +105,28 @@ pub(crate) fn known_defs() -> KnownDefs { KNOWN_DEFS.with(|cell| cell.get()) } +/// Replace the emitted-function set for the current pipeline run. +/// +/// Like [`set_known_defs`], this replaces rather than extends previous +/// thread-local state so one rustc invocation cannot leak into the next. +pub(crate) fn set_emitted_functions(names: std::collections::BTreeSet) { + EMITTED_FUNCTIONS.with(|set| *set.borrow_mut() = names); +} + +/// Whether this exact monomorphized Rust instance will have a `mir.func` +/// definition in the module being translated. +pub(crate) fn is_emitted_function(instance: &mir::mono::Instance) -> bool { + is_emitted_name(&instance.mangled_name()) +} + +/// Whether `mangled_name`, rustc's linkage identity of one monomorphized +/// function, is in the current run's emitted-function set. The name-based +/// form lets the set's replace-not-merge contract be tested without a rustc +/// session. +pub(crate) fn is_emitted_name(mangled_name: &str) -> bool { + EMITTED_FUNCTIONS.with(|set| set.borrow().contains(mangled_name)) +} + /// Validate the source reference and the fully specialized read-only storage /// contract using rustc's type system. MIR locals have erased regions, so the /// outer lifetime is read from the declaration before instantiation instead. @@ -389,6 +417,15 @@ mod pointer_origin { } } + /// ABI RULE: `cuda_device::SharedPtr` stores a `*mut SharedSlot`, so + /// the shared pointer it lowers to is always `RawMut`. + pub(crate) fn abi_shared_ptr_field() -> PointerOrigin { + PointerOrigin { + kind: MirPointerKind::RawMut, + mutable: true, + } + } + /// ABI RULE: `cuda_device::DisjointSlice<'_, T>` stores its data pointer /// as `*mut T`, so its data/element pointers are always `RawMut`. pub(crate) fn abi_disjoint_slice_data_ptr() -> PointerOrigin { @@ -446,60 +483,120 @@ pub(crate) use pointer_origin::*; /// Validity facts that are safe to expose on a Rust-reference kernel parameter. /// -/// Presence of this fact proves non-nullness. `pointee_alignment` is rustc's -/// ABI alignment for the pointee represented by the physical data pointer. -/// No aliasing, lifetime, or dereferenceability guarantee is encoded here. +/// Presence proves non-nullness. Alignment comes from the pointee layout; +/// aliasing and dereferenceability come only from rustc's argument ABI. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct ReferenceParamValidity { pub(crate) pointee_alignment: u64, + pub(crate) no_alias: bool, + pub(crate) dereferenceable_bytes: u64, } -/// Return validity facts for the currently audited Rust-reference kernel scope. -/// -/// The only accepted evidence is the typed `rustc_public` type/layout API: +/// Return typed rustc facts for a source Rust-ABI kernel reference argument. /// -/// * `RigidTy::Ref` proves that the source value is a Rust reference and thus -/// that its data pointer is non-null. -/// * a sized pointee's `Ty::layout()` supplies its ABI alignment; -/// * for `&[T]` / `&mut [T]`, the physical data pointer points at `T`, so the -/// element layout supplies the alignment. +/// Match the source index before physical ABI flattening: ignored ZST arguments +/// still occupy an `FnAbi` position. A sized reference uses its direct pointer +/// attributes; a slice uses the first member of its scalar pair. In particular, +/// an indirect by-value aggregate's carrier attributes never describe a Rust +/// reference or its backing allocation. /// -/// Other DSTs, raw pointers, ADTs such as `DisjointSlice`, and any type whose -/// required typed layout is unavailable fail closed with no fact. +/// `RigidTy::Ref` proves non-nullness. Sized pointee layouts and slice element +/// layouts supply alignment. Other DSTs, raw pointers, and by-value ADTs such as +/// `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, -) -> Option { + fn_abi: &rustc_public::abi::FnAbi, + source_index: usize, +) -> Result, String> { + use rustc_public::abi::PassMode; use rustc_public::ty::{RigidTy, TyKind}; - // Reuse #1186's typed pointer-origin oracle rather than independently - // recognizing reference provenance here. Raw pointers are therefore - // rejected by the same rustc_public-derived classification that mints - // concrete MIR pointer kinds. - let (pointee, origin) = pointer_origin_of_ty(ty)?; + let argument = fn_abi.args.get(source_index).ok_or_else(|| { + format!("kernel source argument index {source_index} is absent from rustc FnAbi") + })?; + if argument.ty != *ty { + return Err(format!( + "kernel source argument {source_index} type does not match rustc FnAbi" + )); + } + + // Reuse the typed pointer-origin oracle so reference provenance is never + // independently inferred from an ABI attribute or pointer-shaped carrier. + let Some((pointee, origin)) = pointer_origin_of_ty(ty) else { + return Ok(None); + }; if !origin.is_reference() { - return None; + return Ok(None); } - let alignment = match pointee.kind() { - TyKind::RigidTy(RigidTy::Slice(element)) => { - let shape = element.layout().ok()?.shape(); - if !shape.is_sized() { - return None; - } - shape.abi_align - } + let (layout_ty, is_slice) = match pointee.kind() { + TyKind::RigidTy(RigidTy::Slice(element)) => (element, true), + _ => (pointee, false), + }; + let Ok(layout) = layout_ty.layout() else { + return Ok(None); + }; + let shape = layout.shape(); + if !shape.is_sized() || shape.abi_align == 0 || !shape.abi_align.is_power_of_two() { + return Ok(None); + } + + match (&argument.mode, is_slice) { + (PassMode::Direct(_), false) | (PassMode::Pair(_, _), true) => {} _ => { - let shape = pointee.layout().ok()?.shape(); - if !shape.is_sized() { - return None; - } - shape.abi_align + return Err(format!( + "kernel reference argument {source_index} has an unsupported rustc FnAbi pass mode" + )); } - }; + } - (alignment != 0 && alignment.is_power_of_two()).then_some(ReferenceParamValidity { - pointee_alignment: alignment, - }) + 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 + // body instead; volatile accesses prevent that inference. + // + // LLVM before 23 (including legacy libNVVM consumers) treats parameter + // 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 no_free { pointee_size_bytes } else { 0 }; + Ok(Some(ReferenceParamValidity { + pointee_alignment: shape.abi_align, + no_alias, + dereferenceable_bytes, + })) } // ============================================================================ @@ -880,8 +977,9 @@ pub fn kernel_caller(value: &u32) { another_kernel(value); } ); } - #[test] - fn reference_param_validity_uses_only_typed_rustc_public_facts() { + fn reference_validity_fixture( + opt_level: u8, + ) -> (Vec>, Vec) { let unique = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .expect("system time before unix epoch") @@ -896,6 +994,9 @@ pub fn kernel_caller(value: &u32) { another_kernel(value); } std::fs::write( &fixture, r#" +use core::cell::UnsafeCell; +use core::marker::PhantomPinned; + #[repr(align(16))] pub struct AlignedZst; @@ -903,7 +1004,18 @@ pub struct ByValue { pub pointer: *const f32, } +pub struct LargeByValue { + pub values: [u64; 8], +} + +pub struct Pinned { + pub value: u32, + pub marker: PhantomPinned, +} + +#[track_caller] pub fn reference_validity( + _ignored: AlignedZst, _shared: &f32, _unique: &mut f32, _slice: &[f32], @@ -912,6 +1024,13 @@ pub fn reference_validity( _zst: &AlignedZst, _raw: *const f32, _by_value: ByValue, + _cell: &UnsafeCell, + _mutable_cell: &mut UnsafeCell, + _pinned: &mut Pinned, + _pinned_shared: &Pinned, + _cell_slice: &[UnsafeCell], + _pinned_slice: &mut [Pinned], + _indirect: LargeByValue, ) {} "#, ) @@ -935,30 +1054,65 @@ pub fn reference_validity( "--crate-name=reference_validity_fixture".to_string(), "--emit=metadata".to_string(), "-Zmir-opt-level=0".to_string(), + format!("-Copt-level={opt_level}"), format!("--out-dir={}", root.display()), format!("--sysroot={sysroot}"), fixture.display().to_string(), ]; - let facts = std::thread::Builder::new() + let result = std::thread::Builder::new() .stack_size(64 * 1024 * 1024) .spawn(move || { rustc_public::run!(&args, || { - use rustc_public::CrateDef; + use rustc_public::abi::PassMode; + use rustc_public::mir::mono; - let body = rustc_public::all_local_items() + let item = rustc_public::all_local_items() .into_iter() .find(|item| item.name().ends_with("::reference_validity")) - .and_then(|item| item.body()) - .expect("fixture function body"); - let facts = body - .locals() + .expect("fixture function"); + let instance = mono::Instance::try_from(item).expect("fixture instance"); + let body = instance.body().expect("fixture function body"); + let fn_abi = instance.fn_abi().expect("fixture function ABI"); + let source_types = body.locals()[1..17] .iter() - .skip(1) - .take(8) - .map(|decl| reference_param_validity(&decl.ty)) + .map(|decl| decl.ty) .collect::>(); - std::ops::ControlFlow::<(), _>::Continue(facts) + let facts = source_types + .iter() + .enumerate() + .map(|(index, ty)| { + reference_param_validity(&instance, ty, &fn_abi, index) + .expect("source parameter matches ABI") + }) + .collect::>(); + + // 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(&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(&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(&instance, &source_types[3], &wrong_mode, 3) + .unwrap_err(), + ); + std::ops::ControlFlow::<(), _>::Continue((facts, errors)) }) }) .unwrap() @@ -967,31 +1121,82 @@ pub fn reference_validity( .expect("in-process fixture compilation succeeds"); std::fs::remove_dir_all(&root).ok(); + result + } + #[test] + fn reference_param_validity_consumes_typed_rustc_abi_facts() { + let (facts, errors) = reference_validity_fixture(3); + let observed = facts + .into_iter() + .map(|fact| { + fact.map(|fact| { + ( + fact.pointee_alignment, + fact.no_alias, + fact.dereferenceable_bytes, + ) + }) + }) + .collect::>(); assert_eq!( - facts, + observed, vec![ - Some(ReferenceParamValidity { - pointee_alignment: 4, - }), - Some(ReferenceParamValidity { - pointee_alignment: 4, - }), - Some(ReferenceParamValidity { - pointee_alignment: 4, - }), - Some(ReferenceParamValidity { - pointee_alignment: 4, - }), - Some(ReferenceParamValidity { - pointee_alignment: 1, - }), - Some(ReferenceParamValidity { - pointee_alignment: 16, - }), - None, - None, + None, // ignored by-value ZST + Some((4, true, 4)), // &f32 + Some((4, true, 4)), // &mut f32 + Some((4, true, 0)), // &[f32]: unknown length + Some((4, true, 0)), // &mut [f32] + Some((1, true, 1)), // &u8 + Some((16, true, 0)), // aligned reference to a ZST + None, // raw pointer + None, // direct by-value ADT + Some((4, false, 0)), // &UnsafeCell + Some((4, true, 4)), // &mut UnsafeCell + Some((4, false, 0)), // &mut !Unpin + Some((4, true, 4)), // shared frozen !Unpin + Some((4, false, 0)), // &[UnsafeCell] + Some((4, false, 0)), // &mut [!Unpin] + None, // indirect by-value carrier ] ); + assert!(errors[0].contains("type"), "{}", errors[0]); + assert!(errors[1].contains("index"), "{}", errors[1]); + assert!(errors[2].contains("pass mode"), "{}", errors[2]); + assert!(errors[3].contains("pass mode"), "{}", errors[3]); + } + + #[test] + fn reference_param_validity_retains_conservative_unoptimized_abi_facts() { + let (facts, _) = reference_validity_fixture(0); + assert_eq!(facts.iter().flatten().count(), 12); + for fact in facts.into_iter().flatten() { + assert!(!fact.no_alias); + assert_eq!(fact.dereferenceable_bytes, 0); + assert!(fact.pointee_alignment.is_power_of_two()); + } + } + + /// `set_emitted_functions` replaces the previous run's set instead of + /// merging into it. A name left over from an earlier rustc invocation on + /// the same thread would keep a diverging call to a function this run + /// never emits, and the module would reference an undefined symbol. + #[test] + fn emitted_function_set_is_replaced_per_run() { + let first_run = std::collections::BTreeSet::from(["first_run_die".to_string()]); + set_emitted_functions(first_run); + assert!(is_emitted_name("first_run_die")); + assert!(!is_emitted_name("second_run_spin")); + + let second_run = std::collections::BTreeSet::from(["second_run_spin".to_string()]); + set_emitted_functions(second_run); + assert!(is_emitted_name("second_run_spin")); + assert!( + !is_emitted_name("first_run_die"), + "an entry from the previous run must not survive the replacement" + ); + + set_emitted_functions(std::collections::BTreeSet::new()); + assert!(!is_emitted_name("second_run_spin")); } } 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/payload_store.rs b/cuda-oxide/crates/mir-importer/src/translator/payload_store.rs index dc892a5f30..3c9482d32d 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/payload_store.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/payload_store.rs @@ -45,7 +45,7 @@ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::TypeHandle; diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/aggregate.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/aggregate.rs index 165fba12bb..18106da4cb 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/aggregate.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/aggregate.rs @@ -22,7 +22,7 @@ use dialect_mir::ops::{MirInsertFieldOp, MirUndefOp}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_error; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::printable::Printable; @@ -589,6 +589,30 @@ pub(super) fn translate_aggregate_rvalue( })?; let tuple_ty = types::translate_type(ctx, &rust_tuple_ty)?; + // Normalize pointer representations to the tuple's declared + // element types before constructing the aggregate. + let expected_types = { + let tuple_ref = tuple_ty.deref(ctx); + tuple_ref + .downcast_ref::() + .expect("translated tuple must have MirTupleType") + .get_types() + .to_vec() + }; + + for (value, expected_type) in element_values.iter_mut().zip(expected_types) { + let (normalized, prev_after_cast) = cast_to_expected_pointer_type_if_needed( + ctx, + *value, + expected_type, + block_ptr, + current_prev_op, + loc.clone(), + ); + *value = normalized; + current_prev_op = prev_after_cast; + } + // Create mir.construct_tuple operation use dialect_mir::ops::MirConstructTupleOp; diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/coerce.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/coerce.rs index d9a3b4913d..b73c154261 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/coerce.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/coerce.rs @@ -12,7 +12,7 @@ use dialect_mir::attributes::{MirCastKindAttr, MirPointerKindAuthorityAttr}; use dialect_mir::ops::MirCastOp; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; @@ -342,6 +342,48 @@ mod tests { use dialect_mir::types::{MirPointerKind, MirPtrType}; use pliron::builtin::types::{IntegerType, Signedness}; + #[test] + fn expected_pointer_normalization_preserves_cluster_shared_raw_const() { + let mut ctx = Context::new(); + crate::translator::register_dialects(&mut ctx); + + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + + let source_ty: TypeHandle = MirPtrType::get_with_kind( + &mut ctx, + pointee, + false, + dialect_mir::types::address_space::CLUSTER_SHARED, + MirPointerKind::RawConst, + ) + .into(); + + let target_ty: TypeHandle = + MirPtrType::get_generic_with_kind(&mut ctx, pointee, false, MirPointerKind::RawConst) + .into(); + + let block = BasicBlock::new(&mut ctx, None, vec![source_ty]); + let source_value = block.deref(&ctx).get_argument(0); + + let (normalized, last_op) = cast_to_expected_pointer_type_if_needed( + &mut ctx, + source_value, + target_ty, + block, + None, + Location::Unknown, + ); + + assert_eq!(normalized.get_type(&ctx), target_ty); + + let cast_op = last_op.expect("AS7 to generic normalization must insert a pointer cast"); + + assert!( + Operation::get_op::(cast_op, &ctx).is_some(), + "normalization must produce mir.cast" + ); + } + #[test] fn expected_pointer_normalization_rejects_concrete_kind_change() { let mut ctx = Context::new(); 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 fe54340822..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 @@ -23,7 +23,7 @@ use dialect_mir::types::MirFP16Type; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::TypeHandle; @@ -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 8bc2ade2f8..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 @@ -21,7 +21,7 @@ use dialect_mir::types::MirFP16Type; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::TypeHandle; @@ -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/const_enum.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_enum.rs index 2a7b611451..768cd4a56b 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_enum.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_enum.rs @@ -17,7 +17,7 @@ use crate::translator::layout::{enum_tag_offset, enum_variant_field_offsets}; use dialect_mir::ops::MirConstructEnumOp; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::TypeHandle; diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_union.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_union.rs index f8b460aaf1..24c8cb2715 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_union.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/const_union.rs @@ -23,7 +23,7 @@ use dialect_mir::types::MirPointerKind; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::TypeHandle; diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/expr.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/expr.rs index 7ce0caf201..7ce8468d31 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/expr.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/expr.rs @@ -27,7 +27,7 @@ use dialect_mir::ops::{ use pliron::basic_block::BasicBlock; use pliron::builtin::types::{FP32Type, FP64Type, IntegerType}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::Typed; @@ -369,6 +369,17 @@ pub fn translate_rvalue( let result_type = types::translate_type(ctx, ty)?; + // A cast whose source and target lower to the same dialect type + // changes nothing. `isize` and `i64` are that pair on this target: + // core steps an `isize` with `overflowing_add(self as i64, rhs as + // i64)` and casts the sum back. Leaving those casts in the step + // block hides the range `for` from the unroller. Return the operand + // and emit no operation, as the reborrow path in this function does + // when its transmute would be between identical types. + if operand_val.get_type(ctx) == result_type { + return Ok((None, operand_val, prev_op_after_operand)); + } + let op = Operation::new( ctx, MirCastOp::get_concrete_op_info(), 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 7a6e0de1dc..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 @@ -12,7 +12,7 @@ use dialect_mir::ops::MirCastOp; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::utils::apint::APInt; @@ -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/operand.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/operand.rs index 92e1be5a80..9e598103e8 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/operand.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/operand.rs @@ -34,7 +34,7 @@ use dialect_mir::types::MirFP16Type; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{FP32Type, FP64Type, IntegerType}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::TypeHandle; 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 0b09daa48c..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 @@ -22,7 +22,7 @@ use dialect_mir::ops::{ use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::printable::Printable; @@ -64,6 +64,7 @@ pub(crate) fn translate_place_address( prev_op: Option>, loc: Location, ) -> TranslationResult>)>> { + let place = &*super::place_read::strip_shared_ptr_field_projections(body, place); let Some(slot) = value_map.get_slot(place.local) else { return Ok(None); }; @@ -171,9 +172,10 @@ pub(super) fn emit_slice_len_extract( /// `ConstantIndex { offset, from_end: true }` on a runtime-length slice. /// /// rustc defines the offset as 1-based from the end, so the index is -/// `slice_len - offset`. The MIR pattern-length test dominates this place, -/// therefore the subtraction cannot underflow on an executed path. -fn emit_from_end_slice_index( +/// `slice_len - offset`. rustc's MIR validator rejects a from-end index whose +/// `offset` exceeds `min_length`, and the MIR pattern-length test dominates +/// this place, therefore the subtraction cannot underflow on an executed path. +pub(super) fn emit_from_end_slice_index( ctx: &mut Context, slice_len: Value, offset: u64, @@ -570,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. @@ -593,11 +595,14 @@ pub(super) fn emit_array_subslice_value( /// out an address that cannot honor the storage coercion. /// `Subslice` is handled for both sized arrays and slice fat pointers. Array /// subslices keep an in-memory address; slice subslices advance the data -/// pointer and rebuild metadata as `len - (from + to)`, and can continue into -/// element Index/forward-ConstantIndex projections. From-end `ConstantIndex` -/// is accepted only when the immediately preceding fat-slice deref supplied -/// runtime length metadata; other bases (including a from-end index after a -/// Subslice) punt instead of guessing a length. +/// pointer and rebuild metadata as `len - (from + to)`. Element Index and +/// ConstantIndex projections can continue from that rebuilt slice. A from-end +/// ConstantIndex uses the rebuilt slice's runtime length metadata, so its index +/// is computed relative to the subslice rather than the original fat slice. +/// Two cases still return `Ok(None)` on purpose: a slice `Subslice` that does +/// not directly follow a fat Deref (a nested `Subslice`, or a `Subslice` after +/// a slice-tail `Field`), and a from-end `ConstantIndex` on a base that +/// carries no slice length. /// /// Returns `Ok(Some((addr, last_op)))` on success, `Ok(None)` if the /// projection chain contains an element this helper doesn't know how to @@ -779,9 +784,7 @@ fn translate_place_addr_from_slot( match &projection[proj_idx + 2] { mir::ProjectionElem::Index(_) - | mir::ProjectionElem::ConstantIndex { - from_end: false, .. - } => { + | mir::ProjectionElem::ConstantIndex { .. } => { let Some(data_ptr_ty) = erased_slice_data_pointer_type(ctx, subslice, elem_ty) else { @@ -804,6 +807,22 @@ fn translate_place_addr_from_slot( current = extract_ptr.deref(ctx).get_result(0); current_prev_op = Some(extract_ptr); current_is_slice_data = true; + current_slice_len = if matches!( + &projection[proj_idx + 2], + mir::ProjectionElem::ConstantIndex { from_end: true, .. } + ) { + let (len, len_op) = emit_slice_len_extract( + ctx, + subslice, + block_ptr, + current_prev_op, + loc.clone(), + ); + current_prev_op = Some(len_op); + Some(len) + } else { + None + }; consumed_slice_subslice = true; continue; } @@ -1342,6 +1361,7 @@ fn translate_place_addr_from_slot( if *from_end && consumed_slice_subslice { consumed_slice_subslice = false; current_is_slice_data = true; + current_slice_len = entered_slice_len; continue; } // A slice Subslice must be lowered while the fat-pointer diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_iter.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_iter.rs index 2beb14d9cd..eb345e6aba 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_iter.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/place_iter.rs @@ -8,7 +8,8 @@ use super::const_bytes::translate_zero_sized_constant_value; use super::const_enum::create_ghost_enum_default; use super::place_addr::{ - emit_array_subslice_value, emit_slice_len_extract, emit_slice_subslice_value, + emit_array_subslice_value, emit_from_end_slice_index, emit_slice_len_extract, + emit_slice_subslice_value, }; use super::place_read::{ apply_deref_projection, apply_enum_field_projection, apply_field_addr_and_load, @@ -23,7 +24,7 @@ use dialect_mir::ops::{MirExtractFieldOp, MirLoadOp, MirPtrOffsetOp}; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::printable::Printable; @@ -536,14 +537,44 @@ pub fn translate_place_iterative( ) ); } - if *from_end { - return input_err!( - loc, - TranslationErr::unsupported( - "ConstantIndex with from_end=true not yet supported" - ) + let runtime_index = if *from_end { + // A from-end ConstantIndex needs runtime slice length. + // After a slice Subslice, current_value is intentionally + // kept as the rebuilt `(data_ptr, new_len)` fat value. + let current_ty = current_value.get_type(ctx); + if !current_ty + .deref(ctx) + .is::() + { + return input_err!( + loc, + TranslationErr::unsupported( + "from-end ConstantIndex requires preserved slice metadata" + .to_string() + ) + ); + } + + let (slice_len, len_op) = emit_slice_len_extract( + ctx, + current_value, + block_ptr, + current_prev_op, + loc.clone(), ); - } + let (index_value, index_op) = emit_from_end_slice_index( + ctx, + slice_len, + *offset, + block_ptr, + Some(len_op), + loc.clone(), + )?; + current_prev_op = Some(index_op); + Some(index_value) + } else { + None + }; let index = *offset as usize; // A projected unsized slice tail is represented as a fat @@ -593,6 +624,14 @@ pub fn translate_place_iterative( match kind { Ok(ConstIndexKind::Array { element_ty }) => { + // A from-end index never reaches this branch: the + // from_end path above requires a MirSliceType, which + // normalizes to its data pointer and takes the Ptr + // branch. The raw offset is a forward field index here. + debug_assert!( + runtime_index.is_none(), + "from-end ConstantIndex reached the array value path" + ); let op = Operation::new( ctx, MirExtractFieldOp::get_concrete_op_info(), @@ -618,30 +657,35 @@ pub fn translate_place_iterative( current_prev_op = Some(extract_op.get_operation()); } Ok(ConstIndexKind::Ptr { element_ty, ptr_ty }) => { - // Create constant index value - let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); - let apint = APInt::from_u32(index as u32, NonZeroUsize::new(32).unwrap()); - let index_attr = - pliron::builtin::attributes::IntegerAttr::new(i32_ty, apint); - use dialect_mir::ops::MirConstantOp; - let const_op = Operation::new( - ctx, - MirConstantOp::get_concrete_op_info(), - vec![i32_ty.into()], - vec![], - vec![], - 0, - ); - const_op.deref_mut(ctx).set_loc(loc.clone()); - let const_mir = MirConstantOp::new(const_op); - const_mir.set_attr_value(ctx, index_attr); - if let Some(prev) = current_prev_op { - const_mir.get_operation().insert_after(ctx, prev); + let index_value = if let Some(index_value) = runtime_index { + index_value } else { - const_mir.get_operation().insert_at_front(block_ptr, ctx); - } - current_prev_op = Some(const_mir.get_operation()); - let index_value = const_mir.get_operation().deref(ctx).get_result(0); + // Create constant index value. + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let apint = + APInt::from_u32(index as u32, NonZeroUsize::new(32).unwrap()); + let index_attr = + pliron::builtin::attributes::IntegerAttr::new(i32_ty, apint); + use dialect_mir::ops::MirConstantOp; + let const_op = Operation::new( + ctx, + MirConstantOp::get_concrete_op_info(), + vec![i32_ty.into()], + vec![], + vec![], + 0, + ); + const_op.deref_mut(ctx).set_loc(loc.clone()); + let const_mir = MirConstantOp::new(const_op); + const_mir.set_attr_value(ctx, index_attr); + if let Some(prev) = current_prev_op { + const_mir.get_operation().insert_after(ctx, prev); + } else { + const_mir.get_operation().insert_at_front(block_ptr, ctx); + } + current_prev_op = Some(const_mir.get_operation()); + const_mir.get_operation().deref(ctx).get_result(0) + }; // Pointer offset let offset_op = Operation::new( @@ -754,9 +798,13 @@ pub fn translate_place_iterative( current_value = subslice; current_prev_op = Some(last_op); - // A subsequent element projection needs only the adjusted - // data pointer. Keep a terminal Subslice as a fat value so - // PtrMetadata/borrow users can observe the rebuilt length. + // A following Index or forward ConstantIndex needs only the + // adjusted data pointer. Keep a terminal Subslice as a fat + // value so PtrMetadata/borrow users can observe the rebuilt + // length, and keep it for a following from-end + // ConstantIndex as well: that arm reads field 1 for the + // length. The filter below must therefore stay restricted + // to Index and forward ConstantIndex. if matches!( place.projection.get(proj_idx + 1), Some(ProjectionElem::Index(_)) 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 d1e6a0f392..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 @@ -23,7 +23,7 @@ use dialect_mir::ops::{MirExtractFieldOp, MirLoadOp, MirUndefOp}; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::printable::Printable; @@ -80,6 +80,12 @@ pub fn translate_place( prev_op: Option>, loc: Location, ) -> TranslationResult<(Value, Option>)> { + let place = &*strip_shared_ptr_field_projections(body, place); + if place.projection.is_empty() + && let Some(value) = value_map.get_direct_value(place.local) + { + return Ok((value, prev_op)); + } match classify_place_read_strategy(ctx, place, value_map)? { PlaceReadStrategy::Address => { if let Some((value, last_op)) = translate_place_load_from_address( @@ -109,6 +115,59 @@ pub fn translate_place( } } +/// Drop the projections of a place onto the field of a `cuda_device::SharedPtr`. +/// +/// `SharedPtr { ptr: *mut SharedSlot }` is a scalar-lowered newtype: its +/// value is its one field, a CTA-shared pointer (see `types::translate_type`). +/// Projecting that field is therefore the identity. Without this, the place +/// walkers would read the pointer value of a `SharedPtr` nested in another +/// place, such as `(region.base).ptr`, as the address of an aggregate. +/// +/// `ThreadIndex`, the other scalar-lowered newtype, needs no strip. Its scalar +/// is an integer, and `dialect-mir` and `mir-lower` already treat +/// `extract_field` on an integer scalar as the identity. The iterative place +/// walker in `place_iter.rs` takes any pointer-typed current value as an +/// address to load from, so only a newtype whose scalar is itself a pointer is +/// misread. Follow-up: have that walker track whether its current value is an +/// address or a value, then drop this strip. +pub(crate) fn strip_shared_ptr_field_projections<'a>( + body: &mir::Body, + place: &'a mir::Place, +) -> std::borrow::Cow<'a, mir::Place> { + use rustc_public::ty::{RigidTy, TyKind}; + + if !place + .projection + .iter() + .any(|elem| matches!(elem, mir::ProjectionElem::Field(0, _))) + { + return std::borrow::Cow::Borrowed(place); + } + let mut place_ty = body.locals()[place.local].ty; + let mut kept = Vec::with_capacity(place.projection.len()); + for elem in &place.projection { + let is_shared_ptr = matches!( + place_ty.kind(), + TyKind::RigidTy(RigidTy::Adt(adt_def, _)) if facts::is_cuda_device_adt(&adt_def, "SharedPtr") + ); + let Ok(next_ty) = elem.ty(place_ty) else { + return std::borrow::Cow::Borrowed(place); + }; + if !(is_shared_ptr && matches!(elem, mir::ProjectionElem::Field(0, _))) { + kept.push(elem.clone()); + } + place_ty = next_ty; + } + if kept.len() == place.projection.len() { + std::borrow::Cow::Borrowed(place) + } else { + std::borrow::Cow::Owned(mir::Place { + local: place.local, + projection: kept, + }) + } +} + // ============================================================================ // Place Read Strategy and Address Path // ============================================================================ @@ -716,10 +775,16 @@ fn translate_place_value_fallback( // back into SSA extracts for promotable arrays. let index = if *from_end { + // rustc emits `from_end: true` only for slice patterns + // (`from_end: !is_array` in match_pair.rs), and its MIR + // validator rejects an array indexed from the end. A + // bare local is sized, so valid MIR never reaches this + // arm; it only reports a malformed input. return input_err!( loc, TranslationErr::unsupported( - "ConstantIndex with from_end=true not yet supported" + "from-end ConstantIndex on an array local is not valid MIR \ + (rustc emits from_end only for slice patterns)" ) ); } else { diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/pointee.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/pointee.rs index f92a13e1c5..1e2f1549e7 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/pointee.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/pointee.rs @@ -9,7 +9,7 @@ use crate::translator::facts; use dialect_mir::ops::{MirExtractFieldOp, MirPtrOffsetOp}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/promoted.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/promoted.rs index 2a5a0eeaa6..db557c7999 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/promoted.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/promoted.rs @@ -25,7 +25,7 @@ use dialect_mir::types::MirFP16Type; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::TypeHandle; diff --git a/cuda-oxide/crates/mir-importer/src/translator/rvalue/static_global.rs b/cuda-oxide/crates/mir-importer/src/translator/rvalue/static_global.rs index 14caa35eba..0bb528ae46 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/rvalue/static_global.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/rvalue/static_global.rs @@ -20,7 +20,7 @@ use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; diff --git a/cuda-oxide/crates/mir-importer/src/translator/statement.rs b/cuda-oxide/crates/mir-importer/src/translator/statement.rs index 69569ea4b7..74588c3098 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/statement.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/statement.rs @@ -48,7 +48,7 @@ use dialect_mir::types::MirEnumType; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::printable::Printable; @@ -304,10 +304,17 @@ pub fn translate_statement( // local's slot and emit `mir.store`. if *from_end { + // rustc emits `from_end: true` only for slice + // patterns (`from_end: !is_array` in match_pair.rs), + // and its MIR validator rejects an array indexed + // from the end. A bare local is sized, so valid MIR + // never reaches this arm; it only reports a + // malformed input. return input_err!( loc, TranslationErr::unsupported( - "ConstantIndex with from_end=true not yet supported for writes" + "from-end ConstantIndex store on an array local is not valid \ + MIR (rustc emits from_end only for slice patterns)" ) ); } diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/drop_glue.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/drop_glue.rs index 9ff962ad68..3d95fc512e 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/drop_glue.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/drop_glue.rs @@ -8,8 +8,8 @@ //! This module provides two capabilities: //! //! 1. **No-op analysis** ([`drop_glue_is_noop`]): a conservative static proof -//! that dropping a value does nothing observable, allowing the `Drop` -//! terminator to be lowered as a plain branch (fast path). +//! that dropping a value does nothing observable and returns, allowing +//! the `Drop` terminator to be lowered as a plain branch (fast path). //! //! 2. **Drop glue emission** (`emit_drop_glue`): when the no-op proof fails, //! emits a device-side call to the monomorphized `drop_in_place::` @@ -42,6 +42,11 @@ //! Anything else (writes through pointers, asserts, inline asm, calls we //! cannot resolve or whose body we cannot see) fails the proof. //! +//! The proof also requires the glue to return. A reachable control-flow +//! cycle, or a call back into an instance that is already under proof, +//! fails it even when every statement is harmless: erasing a call that +//! may never return would let the code after the drop execute. +//! //! The walk follows constant `switchInt` discriminants, so branches that //! the compiler has already folded shut (for example checks behind the //! `UbChecks` flag, which is off for device builds) do not have to be @@ -73,7 +78,7 @@ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::identifier::Legaliser; use pliron::input_error; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::Typed; @@ -320,16 +325,17 @@ fn compute_drop_place_address( } // ============================================================================ -// No-op analysis (unchanged from original) +// Conservative no-op and termination analysis // ============================================================================ -/// Proves that calling `instance` does nothing observable. +/// Proves that erasing a call to `instance` preserves observable behavior, +/// including whether control can return to its caller. /// /// `in_progress` holds the mangled names of instances currently being -/// proven further up the call stack. If we meet one of them again we -/// treat the cycle as harmless: a cycle cannot *introduce* an -/// observable effect, every effect would already have failed the proof -/// on some statement or terminator along the way. +/// proven further up the call stack. Meeting one again means termination +/// is not proven: even recursion with no externally visible writes can +/// diverge, and erasing such a call would make unreachable code execute. +/// Recursion therefore fails this conservative proof. fn instance_is_noop(instance: &Instance, in_progress: &mut Vec) -> bool { // Fast path: a type with no drop glue at all resolves to an "empty" // drop shim that exists only to fill vtable slots. rustc_public @@ -344,7 +350,7 @@ fn instance_is_noop(instance: &Instance, in_progress: &mut Vec) -> bool let name = instance.mangled_name(); if in_progress.contains(&name) { - return true; + return false; } // No body means we cannot see what the call does (an intrinsic, a @@ -359,102 +365,136 @@ fn instance_is_noop(instance: &Instance, in_progress: &mut Vec) -> bool result } -/// Walks every block of `body` reachable from the entry block and -/// checks that nothing observable happens on the way to `return`. +/// Walks every block of `body` reachable from the entry block and proves +/// that execution is effect-free and contains no reachable control-flow cycle. +/// +/// A cycle fails the proof even when every statement in it is locally harmless: +/// removing a call that can diverge would change whether its continuation runs. fn body_is_noop(body: &mir::Body, in_progress: &mut Vec) -> bool { - let mut visited = vec![false; body.blocks.len()]; - let mut worklist: Vec = vec![0]; + #[derive(Clone, Copy, PartialEq, Eq)] + enum BlockState { + Unseen, + InProgress, + Proven, + } - while let Some(idx) = worklist.pop() { - if std::mem::replace(&mut visited[idx], true) { - continue; - } - let block = &body.blocks[idx]; + #[derive(Clone, Copy)] + enum WorkItem { + Enter(mir::BasicBlockIdx), + Finish(mir::BasicBlockIdx), + } - for stmt in &block.statements { - if !statement_is_noop(&stmt.kind) { - return false; - } - } + let mut states = vec![BlockState::Unseen; body.blocks.len()]; + let mut worklist = vec![WorkItem::Enter(0)]; - match &block.terminator.kind { - // Reaching `return` with only no-op work behind us is the - // success case. `unreachable` cannot execute in a valid - // program, so it cannot contribute an effect either. - mir::TerminatorKind::Return | mir::TerminatorKind::Unreachable => {} - - mir::TerminatorKind::Goto { target } => worklist.push(*target), - - mir::TerminatorKind::SwitchInt { discr, targets } => { - match const_operand_bits(discr) { - // Known discriminant: only the matching branch can - // run, so only that branch needs to be a no-op. - // This is what skips the dead "really drop the - // elements" branch in `IntoIter`'s destructor. - Some(value) => { - let target = targets - .branches() - .find(|(branch_value, _)| *branch_value == value) - .map(|(_, target)| target) - .unwrap_or_else(|| targets.otherwise()); - worklist.push(target); - } - // Unknown discriminant: every branch must be a - // no-op. - None => worklist.extend(targets.all_targets()), - } + while let Some(item) = worklist.pop() { + match item { + WorkItem::Finish(idx) => { + debug_assert!(matches!(states[idx], BlockState::InProgress)); + states[idx] = BlockState::Proven; + continue; } - // A nested drop is fine when the dropped value's own glue - // passes this same proof. - mir::TerminatorKind::Drop { place, target, .. } => { - let Ok(place_ty) = place.ty(body.locals()) else { - return false; - }; - let nested = Instance::resolve_drop_in_place(place_ty); - if !instance_is_noop(&nested, in_progress) { - return false; + WorkItem::Enter(idx) => { + match states[idx] { + // A completed block reached through another predecessor is + // ordinary CFG reconvergence. + BlockState::Proven => continue, + + // A block already on the active DFS path is a reachable CFG + // cycle. Effect-freedom does not prove that it terminates. + BlockState::InProgress => return false, + + BlockState::Unseen => {} } - worklist.push(*target); - } - // A call is fine when we can resolve exactly which function - // runs and that function passes this same proof. The - // `drop_in_place` shim for a type with an `impl Drop` is a - // single such call to `::drop`. - mir::TerminatorKind::Call { - func, - destination, - target: Some(target), - .. - } => { - if place_writes_through_pointer(destination) { - return false; + states[idx] = BlockState::InProgress; + let block = &body.blocks[idx]; + + for stmt in &block.statements { + if !statement_is_noop(&stmt.kind) { + return false; + } } - let Ok(func_ty) = func.ty(body.locals()) else { - return false; - }; - let TyKind::RigidTy(RigidTy::FnDef(def, args)) = func_ty.kind() else { - // A function pointer or other indirect callee: we - // do not know what runs. - return false; - }; - let Ok(callee) = Instance::resolve(def, &args) else { - return false; - }; - if !instance_is_noop(&callee, in_progress) { - return false; + + // Only mark this block Proven after every runtime-reachable + // successor has itself been proven. + worklist.push(WorkItem::Finish(idx)); + + match &block.terminator.kind { + mir::TerminatorKind::Return | mir::TerminatorKind::Unreachable => {} + + mir::TerminatorKind::Goto { target } => { + worklist.push(WorkItem::Enter(*target)); + } + + mir::TerminatorKind::SwitchInt { discr, targets } => { + match const_operand_bits(discr) { + Some(value) => { + let target = targets + .branches() + .find(|(branch_value, _)| *branch_value == value) + .map(|(_, target)| target) + .unwrap_or_else(|| targets.otherwise()); + + worklist.push(WorkItem::Enter(target)); + } + + None => { + for target in targets.all_targets() { + worklist.push(WorkItem::Enter(target)); + } + } + } + } + + mir::TerminatorKind::Drop { place, target, .. } => { + let Ok(place_ty) = place.ty(body.locals()) else { + return false; + }; + + let nested = Instance::resolve_drop_in_place(place_ty); + if !instance_is_noop(&nested, in_progress) { + return false; + } + + worklist.push(WorkItem::Enter(*target)); + } + + mir::TerminatorKind::Call { + func, + destination, + target: Some(target), + .. + } => { + if place_writes_through_pointer(destination) { + return false; + } + + let Ok(func_ty) = func.ty(body.locals()) else { + return false; + }; + + let TyKind::RigidTy(RigidTy::FnDef(def, args)) = func_ty.kind() else { + return false; + }; + + let Ok(callee) = Instance::resolve(def, &args) else { + return false; + }; + + if !instance_is_noop(&callee, in_progress) { + return false; + } + + worklist.push(WorkItem::Enter(*target)); + } + + // Asserts, diverging calls, inline asm, resume/abort, and + // unknown terminators do not prove normal termination. + _ => return false, } - worklist.push(*target); } - - // Everything else (asserts, diverging calls, inline asm, - // resume/abort) either has an effect or might not return. - // Asserts fail the proof deliberately: following only the - // success edge would silently delete a would-be device trap. - // A failed proof now emits the real drop_in_place call, so - // there is no pressure to widen the proof here. - _ => return false, } } 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 d9d45fef48..78012fbfe0 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs @@ -31,7 +31,7 @@ use pliron::builtin::type_interfaces::FunctionTypeInterface; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; @@ -316,10 +316,10 @@ pub fn bundle_generated_u32_results_as_array( /// 1. Translate all MIR arguments to Pliron IR values /// 2. At a foreign ABI boundary, adapt pointer address spaces to the exact /// declared parameter types without changing pointer kind or mutability -/// 3. Prepare the call destination using the current projection state +/// 3. For returning calls, prepare the destination using the current projection state /// 4. Create a `mir.call` operation carrying the callee's name attribute -/// 5. Store the result through the prepared destination -/// 6. Emit a zero-operand goto to the call's success target +/// 5. For returning calls, store the result and goto the success target +/// 6. For diverging calls, emit `mir.unreachable` after the call /// /// Reference arguments (`&mut local`) are handed the local's alloca slot /// pointer directly, so callee writes through the reference are observed by @@ -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, @@ -409,7 +409,7 @@ pub fn emit_function_call( } }; let declared = crate::translator::types::translate_type(ctx, &rust_type)?; - let (normalized, tail) = normalize_foreign_pointer_argument( + let (normalized, tail) = normalize_call_pointer_argument( ctx, *argument, declared, @@ -422,16 +422,21 @@ pub fn emit_function_call( } } - let (prepared_destination, prepared_last_op) = prepare_destination_write( - ctx, - body, - destination, - value_map, - block_ptr, - last_op, - loc.clone(), - )?; - last_op = prepared_last_op; + let prepared_destination = if target.is_some() { + let (prepared, prepared_last_op) = prepare_destination_write( + ctx, + body, + destination, + value_map, + block_ptr, + last_op, + loc.clone(), + )?; + last_op = prepared_last_op; + Some(prepared) + } else { + None + }; use pliron::builtin::attributes::StringAttr; @@ -462,11 +467,20 @@ pub fn emit_function_call( call_op }; + if target.is_none() { + return Ok(super::emit_unreachable_after( + ctx, + block_ptr, + Some(call_op), + loc, + )); + } + let result_value = call_op.deref(ctx).get_result(0); let goto_prev = finish_destination_write( ctx, - prepared_destination, + prepared_destination.expect("returning call prepared its destination"), result_value, value_map, block_ptr, @@ -474,19 +488,18 @@ pub fn emit_function_call( loc.clone(), )?; - if let Some(target_idx) = target { - Ok(emit_goto(ctx, *target_idx, goto_prev, block_map, loc)) - } else { - input_err!( - loc.clone(), - TranslationErr::unsupported("Call terminator without target not supported".to_string(),) - ) - } + Ok(emit_goto( + ctx, + target.expect("returning call has a success target"), + goto_prev, + block_map, + loc, + )) } -/// Adapt a foreign or Rust call pointer argument only when the address space -/// is the sole difference from its declared parameter type. -fn normalize_foreign_pointer_argument( +/// Adapt a call pointer argument only when the address space is the +/// sole difference from its declared ABI parameter type. +fn normalize_call_pointer_argument( ctx: &mut Context, value: Value, declared_type: TypeHandle, @@ -898,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, @@ -939,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, @@ -967,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, @@ -1000,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, @@ -1030,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/asm.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/asm.rs index a9b24a6ab2..e54eb973b7 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/asm.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/asm.rs @@ -16,7 +16,7 @@ use dialect_nvvm::ops::InlinePtxOp; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/atomic.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/atomic.rs index 97ec48393e..98b7fc12ce 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/atomic.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/atomic.rs @@ -78,7 +78,7 @@ use dialect_mir::types::MirFP16Type; use pliron::basic_block::BasicBlock; use pliron::builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::Typed; diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/debug.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/debug.rs index c51a4c1777..19e270dc49 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/debug.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/debug.rs @@ -21,7 +21,7 @@ use dialect_nvvm::ops::{AssertFailOp, VprintfOp}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/float_math.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/float_math.rs index e77e0c21e1..77140b8dcb 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/float_math.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/float_math.rs @@ -174,6 +174,58 @@ pub enum RustFloatMathIntrinsic { } impl RustFloatMathIntrinsic { + /// Resolve both the historical width-specific paths and the generic unary + /// intrinsics introduced by the pinned nightly. Generic calls carry their + /// width in rustc's monomorphized type arguments, never in a display name. + pub fn from_core_path_with_substs( + name: &str, + type_substs: &[rustc_public::ty::Ty], + ) -> Result, String> { + use rustc_public::ty::{RigidTy, TyKind}; + + if let Some(intrinsic) = Self::from_core_path(name) { + return Ok(Some(intrinsic)); + } + let float_ty = match type_substs { + [ty] => match ty.kind() { + TyKind::RigidTy(RigidTy::Float(float_ty)) => Some(float_ty), + _ => None, + }, + _ => None, + }; + Self::from_generic_core_path(name, float_ty) + } + + fn from_generic_core_path( + name: &str, + float_ty: Option, + ) -> Result, String> { + use rustc_public::ty::FloatTy; + + // These exact declarations are `fn(T) -> T`. + // Reuse the existing f32/f64 placeholders and libdevice lowerings; + // unsupported widths must not be silently narrowed. + let (f32_intrinsic, f64_intrinsic) = match name { + "core::intrinsics::sin" | "std::intrinsics::sin" => (Self::SinF32, Self::SinF64), + "core::intrinsics::cos" | "std::intrinsics::cos" => (Self::CosF32, Self::CosF64), + "core::intrinsics::exp" | "std::intrinsics::exp" => (Self::ExpF32, Self::ExpF64), + "core::intrinsics::exp2" | "std::intrinsics::exp2" => (Self::Exp2F32, Self::Exp2F64), + "core::intrinsics::log" | "std::intrinsics::log" => (Self::LogF32, Self::LogF64), + "core::intrinsics::log2" | "std::intrinsics::log2" => (Self::Log2F32, Self::Log2F64), + "core::intrinsics::log10" | "std::intrinsics::log10" => { + (Self::Log10F32, Self::Log10F64) + } + _ => return Ok(None), + }; + match float_ty { + Some(FloatTy::F32) => Ok(Some(f32_intrinsic)), + Some(FloatTy::F64) => Ok(Some(f64_intrinsic)), + _ => Err(format!( + "generic float intrinsic `{name}` requires one f32 or f64 type argument; found {float_ty:?}" + )), + } + } + /// Recognize the libcore intrinsic path that survived into MIR. pub fn from_core_path(name: &str) -> Option { match name { @@ -525,7 +577,6 @@ pub fn emit_sincos( use dialect_mir::ops::{MirCallOp, MirConstructTupleOp}; use pliron::builtin::attributes::StringAttr; use pliron::input_err; - use pliron::location::Located; use pliron::op::Op; // Destination tuple type and its scalar element type. Typed from the @@ -685,6 +736,61 @@ mod tests { use super::*; use dialect_mir::rust_intrinsics; + #[test] + fn generic_float_intrinsics_resolve_only_the_declared_float_type() { + use RustFloatMathIntrinsic as I; + use rustc_public::ty::FloatTy; + for (name, f32_variant, f64_variant) in [ + ("sin", I::SinF32, I::SinF64), + ("cos", I::CosF32, I::CosF64), + ("exp", I::ExpF32, I::ExpF64), + ("exp2", I::Exp2F32, I::Exp2F64), + ("log", I::LogF32, I::LogF64), + ("log2", I::Log2F32, I::Log2F64), + ("log10", I::Log10F32, I::Log10F64), + ] { + for root in ["core", "std"] { + let path = format!("{root}::intrinsics::{name}"); + assert_eq!( + I::from_generic_core_path(&path, Some(FloatTy::F32)).unwrap(), + Some(f32_variant) + ); + assert_eq!( + I::from_generic_core_path(&path, Some(FloatTy::F64)).unwrap(), + Some(f64_variant) + ); + for unsupported in [None, Some(FloatTy::F16), Some(FloatTy::F128)] { + let error = I::from_generic_core_path(&path, unsupported).unwrap_err(); + assert!( + error.contains("requires one f32 or f64 type argument"), + "{error}" + ); + } + } + } + for path in [ + "user::intrinsics::exp", + "user::core::intrinsics::sin", + "core::intrinsics::exp_m1", + "core::intrinsics::sqrt", + "core::intrinsics::expf32", + ] { + assert_eq!( + I::from_generic_core_path(path, Some(FloatTy::F32)).unwrap(), + None, + "{path}" + ); + } + assert_eq!( + I::from_core_path("core::intrinsics::expf32"), + Some(I::ExpF32) + ); + assert_eq!( + I::from_core_path("core::intrinsics::expf64"), + Some(I::ExpF64) + ); + } + /// `f32::max` / `f64::max` / `f32::min` / `f64::min` all lower to the /// `_nsz` flavor of the rustc maxNum/minNum intrinsics. Lock the four /// `core::intrinsics::*` paths and their `std::intrinsics::*` aliases to 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 7363584ce2..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -15,7 +15,7 @@ use crate::translator::{terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::ActiveMaskOp; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 54643ce57c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `clc` intrinsics. @@ -19,7 +19,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 ea162f2edc..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,7 @@ use dialect_nvvm::ops::{ClusterBarrierModeAttr, ClusterBarrierOp}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 695d3c8f6c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -15,7 +15,7 @@ use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::{DsmemReadU32Op, MapaSharedClusterOp}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 017ed0c92e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -21,7 +21,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 2bd7ca2bae..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,7 @@ use dialect_nvvm::ops::{BreakpointOp, PmEventOp, TrapOp}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 14363fb24f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `dotprod` intrinsics. @@ -15,7 +15,7 @@ use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::{Dp2aS32Op, Dp2aU32Op, Dp4aS32Op, Dp4aU32Op}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 7d585a7a52..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `elect` intrinsics. @@ -17,7 +17,7 @@ use dialect_nvvm::ops::ElectSyncOp; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 697372cc9e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -21,7 +21,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 07e0c4f542..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -18,7 +18,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 7e76c72bb5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -20,7 +20,7 @@ use pliron::basic_block::BasicBlock; use pliron::builtin::types::{FP32Type, FP64Type, IntegerType, Signedness}; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; 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 62a9d968a9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -18,7 +18,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 cc0df48deb..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `ldmatrix` intrinsics. @@ -18,7 +18,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 6a719bb8a1..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -19,7 +19,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 48657a9535..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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,16 +13,18 @@ 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; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; @@ -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 c1370c9621..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -313,10 +313,14 @@ pub fn is_generated_intrinsic_path(name: &str) -> bool { | "cuda_device::tma::cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0328" | "cuda_device::tma::cp_async_bulk_tensor_1d_g2s" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034" + | "cuda_device::tma::cp_async_bulk_tensor_1d_g2s_cta" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0335" | "cuda_device::tma::cp_async_bulk_tensor_1d_s2g" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0329" | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035" + | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_cta" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0330" | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_multicast" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0331" @@ -325,14 +329,20 @@ pub fn is_generated_intrinsic_path(name: &str) -> bool { | "cuda_device::tma::cp_async_bulk_tensor_2d_s2g" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0332" | "cuda_device::tma::cp_async_bulk_tensor_3d_g2s" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036" + | "cuda_device::tma::cp_async_bulk_tensor_3d_g2s_cta" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0337" | "cuda_device::tma::cp_async_bulk_tensor_3d_s2g" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0333" | "cuda_device::tma::cp_async_bulk_tensor_4d_g2s" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037" + | "cuda_device::tma::cp_async_bulk_tensor_4d_g2s_cta" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0338" | "cuda_device::tma::cp_async_bulk_tensor_4d_s2g" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0334" | "cuda_device::tma::cp_async_bulk_tensor_5d_g2s" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038" + | "cuda_device::tma::cp_async_bulk_tensor_5d_g2s_cta" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0339" | "cuda_device::tma::cp_async_bulk_tensor_5d_s2g" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0928" @@ -517,18 +527,26 @@ pub fn is_generated_intrinsic_path(name: &str) -> bool { | "cuda_device::convert::cvt_rn_relu_satfinite_tf32_f32" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0371" | "cuda_device::convert::cvt_rn_relu_tf32_f32" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030" + | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0822" | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f16x2" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0259" | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f32" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032" + | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0824" | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f16x2" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0261" | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f32" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031" + | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0823" | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f16x2" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0260" | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f32" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033" + | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0825" | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f16x2" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0262" @@ -601,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" @@ -830,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" @@ -848,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" @@ -1158,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" @@ -2301,16 +2327,21 @@ pub fn is_raw_generated_intrinsic_path(name: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0893" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0922" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0328" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0335" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0329" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0330" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0331" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0336" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0332" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0337" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0333" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0338" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0334" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0339" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0928" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0929" @@ -2403,12 +2434,16 @@ pub fn is_raw_generated_intrinsic_path(name: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0083" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0373" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0371" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0822" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0259" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0824" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0261" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0823" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0260" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0825" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0262" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0372" @@ -2445,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" @@ -2559,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" @@ -2568,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" @@ -2723,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" @@ -3431,10 +3470,14 @@ pub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> { } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0328" | "cuda_device::tma::cp_async_bulk_tensor_1d_g2s" => Some("v1:i0328"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034" + | "cuda_device::tma::cp_async_bulk_tensor_1d_g2s_cta" => Some("v1:i1034"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0335" | "cuda_device::tma::cp_async_bulk_tensor_1d_s2g" => Some("v1:i0335"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0329" | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s" => Some("v1:i0329"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035" + | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_cta" => Some("v1:i1035"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0330" | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_multicast" => Some("v1:i0330"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0331" @@ -3443,14 +3486,20 @@ pub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> { | "cuda_device::tma::cp_async_bulk_tensor_2d_s2g" => Some("v1:i0336"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0332" | "cuda_device::tma::cp_async_bulk_tensor_3d_g2s" => Some("v1:i0332"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036" + | "cuda_device::tma::cp_async_bulk_tensor_3d_g2s_cta" => Some("v1:i1036"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0337" | "cuda_device::tma::cp_async_bulk_tensor_3d_s2g" => Some("v1:i0337"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0333" | "cuda_device::tma::cp_async_bulk_tensor_4d_g2s" => Some("v1:i0333"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037" + | "cuda_device::tma::cp_async_bulk_tensor_4d_g2s_cta" => Some("v1:i1037"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0338" | "cuda_device::tma::cp_async_bulk_tensor_4d_s2g" => Some("v1:i0338"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0334" | "cuda_device::tma::cp_async_bulk_tensor_5d_g2s" => Some("v1:i0334"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038" + | "cuda_device::tma::cp_async_bulk_tensor_5d_g2s_cta" => Some("v1:i1038"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0339" | "cuda_device::tma::cp_async_bulk_tensor_5d_s2g" => Some("v1:i0339"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0928" @@ -3635,18 +3684,26 @@ pub fn generated_intrinsic_marker(name: &str) -> Option<&'static str> { | "cuda_device::convert::cvt_rn_relu_satfinite_tf32_f32" => Some("v1:i0373"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0371" | "cuda_device::convert::cvt_rn_relu_tf32_f32" => Some("v1:i0371"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030" + | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2" => Some("v1:i1030"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0822" | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f16x2" => Some("v1:i0822"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0259" | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f32" => Some("v1:i0259"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032" + | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2" => Some("v1:i1032"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0824" | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f16x2" => Some("v1:i0824"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0261" | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f32" => Some("v1:i0261"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031" + | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2" => Some("v1:i1031"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0823" | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f16x2" => Some("v1:i0823"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0260" | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f32" => Some("v1:i0260"), + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033" + | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2" => Some("v1:i1033"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0825" | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f16x2" => Some("v1:i0825"), "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0262" @@ -3721,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" @@ -3957,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" @@ -3975,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") } @@ -4287,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 271dc99c59..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `movmatrix` intrinsics. @@ -15,7 +15,7 @@ use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::MovmatrixTransB16Op; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 5f22170a73..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -21,7 +21,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 95ae3a7713..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -15,7 +15,7 @@ use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::{PackedAtomicAddOp, PackedAtomicFormatAttr}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 811e9681b3..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -15,14 +15,16 @@ use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::{ CvtF16x2F32Op, CvtF32x2Bf16x2Op, CvtRnF16x2E4m3x2Op, CvtRnF16x2E5m2x2Op, CvtRnReluBf16x2F32Op, CvtRnReluF16x2E4m3x2Op, CvtRnReluF16x2E5m2x2Op, CvtRnReluF16x2F32Op, - CvtRnSatfiniteE4m3x2F16x2Op, CvtRnSatfiniteE4m3x2F32Op, CvtRnSatfiniteE5m2x2F16x2Op, - CvtRnSatfiniteE5m2x2F32Op, CvtRnSatfiniteReluE4m3x2F16x2Op, CvtRnSatfiniteReluE4m3x2F32Op, + CvtRnSatfiniteE4m3x2Bf16x2Op, CvtRnSatfiniteE4m3x2F16x2Op, CvtRnSatfiniteE4m3x2F32Op, + CvtRnSatfiniteE5m2x2Bf16x2Op, CvtRnSatfiniteE5m2x2F16x2Op, CvtRnSatfiniteE5m2x2F32Op, + CvtRnSatfiniteReluE4m3x2Bf16x2Op, CvtRnSatfiniteReluE4m3x2F16x2Op, + CvtRnSatfiniteReluE4m3x2F32Op, CvtRnSatfiniteReluE5m2x2Bf16x2Op, CvtRnSatfiniteReluE5m2x2F16x2Op, CvtRnSatfiniteReluE5m2x2F32Op, CvtRzBf16x2F32Op, CvtRzF16x2F32Op, }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; @@ -389,6 +391,45 @@ pub(super) fn try_dispatch( "cvt_rn_relu_f16x2_f32 call without target block", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030" + | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2" => { + require_arity(name, args.len(), 1, &loc)?; + let (arg0, last_op) = rvalue::translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let intrinsic = CvtRnSatfiniteE4m3x2Bf16x2Op::build(ctx, arg0); + intrinsic.deref_mut(ctx).set_loc(loc.clone()); + helpers::set_generated_intrinsic_marker(ctx, intrinsic, "v1:i1030"); + 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, intrinsic, block_ptr, prepared_last_op); + let result = intrinsic.deref(ctx).get_result(0); + Ok(Some(helpers::emit_prepared_result_and_goto( + ctx, + prepared_destination, + result, + target, + block_ptr, + intrinsic, + value_map, + block_map, + loc, + "cvt_rn_satfinite_e4m3x2_bf16x2 call without target block", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0822" | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f16x2" => { require_arity(name, args.len(), 1, &loc)?; @@ -476,6 +517,45 @@ pub(super) fn try_dispatch( "cvt_rn_satfinite_e4m3x2_f32 call without target block", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032" + | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2" => { + require_arity(name, args.len(), 1, &loc)?; + let (arg0, last_op) = rvalue::translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let intrinsic = CvtRnSatfiniteE5m2x2Bf16x2Op::build(ctx, arg0); + intrinsic.deref_mut(ctx).set_loc(loc.clone()); + helpers::set_generated_intrinsic_marker(ctx, intrinsic, "v1:i1032"); + 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, intrinsic, block_ptr, prepared_last_op); + let result = intrinsic.deref(ctx).get_result(0); + Ok(Some(helpers::emit_prepared_result_and_goto( + ctx, + prepared_destination, + result, + target, + block_ptr, + intrinsic, + value_map, + block_map, + loc, + "cvt_rn_satfinite_e5m2x2_bf16x2 call without target block", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0824" | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f16x2" => { require_arity(name, args.len(), 1, &loc)?; @@ -563,6 +643,45 @@ pub(super) fn try_dispatch( "cvt_rn_satfinite_e5m2x2_f32 call without target block", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031" + | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2" => { + require_arity(name, args.len(), 1, &loc)?; + let (arg0, last_op) = rvalue::translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let intrinsic = CvtRnSatfiniteReluE4m3x2Bf16x2Op::build(ctx, arg0); + intrinsic.deref_mut(ctx).set_loc(loc.clone()); + helpers::set_generated_intrinsic_marker(ctx, intrinsic, "v1:i1031"); + 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, intrinsic, block_ptr, prepared_last_op); + let result = intrinsic.deref(ctx).get_result(0); + Ok(Some(helpers::emit_prepared_result_and_goto( + ctx, + prepared_destination, + result, + target, + block_ptr, + intrinsic, + value_map, + block_map, + loc, + "cvt_rn_satfinite_relu_e4m3x2_bf16x2 call without target block", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0823" | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f16x2" => { require_arity(name, args.len(), 1, &loc)?; @@ -650,6 +769,45 @@ pub(super) fn try_dispatch( "cvt_rn_satfinite_relu_e4m3x2_f32 call without target block", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033" + | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2" => { + require_arity(name, args.len(), 1, &loc)?; + let (arg0, last_op) = rvalue::translate_operand( + ctx, + body, + &args[0], + value_map, + block_ptr, + prev_op, + loc.clone(), + )?; + let intrinsic = CvtRnSatfiniteReluE5m2x2Bf16x2Op::build(ctx, arg0); + intrinsic.deref_mut(ctx).set_loc(loc.clone()); + helpers::set_generated_intrinsic_marker(ctx, intrinsic, "v1:i1033"); + 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, intrinsic, block_ptr, prepared_last_op); + let result = intrinsic.deref(ctx).get_result(0); + Ok(Some(helpers::emit_prepared_result_and_goto( + ctx, + prepared_destination, + result, + target, + block_ptr, + intrinsic, + value_map, + block_map, + loc, + "cvt_rn_satfinite_relu_e5m2x2_bf16x2 call without target block", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0825" | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f16x2" => { require_arity(name, args.len(), 1, &loc)?; 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 47f3018ae3..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `prmt` intrinsics. @@ -15,7 +15,7 @@ use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::{PrmtModeAttr, PrmtOp}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 acb9d2d2be..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `redux` intrinsics. @@ -20,7 +20,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 2d333957b1..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -21,7 +21,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; @@ -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 2593523055..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -18,7 +18,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 e91e9a70f9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 e33d6cec27..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -18,7 +18,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 9ad19f8a62..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -22,7 +22,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 0817a6c2da..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6789f53f39..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `stmatrix` intrinsics. @@ -18,7 +18,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 e0d6949b1a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `sync` intrinsics. @@ -16,7 +16,7 @@ use dialect_nvvm::ops::{Barrier0Op, ThreadfenceBlockOp, ThreadfenceOp, Threadfen use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 1f4f183f86..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -32,7 +32,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 ee83464604..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -47,7 +47,7 @@ use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; 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 3c4cff1501..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -19,7 +19,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 02cf1c9317..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -30,7 +30,7 @@ use pliron::basic_block::BasicBlock; use pliron::builtin::types::FP32Type; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; 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 e7123a1cbb..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -42,7 +42,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; 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 299cdaee53..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 04e44ae65b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 522e70ad78..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 b1525803a9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fab4c9c642..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 365e539146..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 57bca785f5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 63a533fa45..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 7b3ca56b96..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 01a5ccca63..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 9e25c77525..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 2cdb7b9a67..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 569197d5fa..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 0670447d85..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: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 cc78b6207b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 d52350e08f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 453444a7f5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c6c75b16f3..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 04c8091720..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `packed_conversion` intrinsics. @@ -300,6 +300,42 @@ fn cvt_rn_relu_f16x2_f32_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn cvt_rn_satfinite_e4m3x2_bf16x2_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030"), + Some("v1:i1030") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::convert::cvt_rn_satfinite_e4m3x2_bf16x2" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2"), + Some("v1:i1030") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1030" + )); +} +#[test] fn cvt_rn_satfinite_e4m3x2_f16x2_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0822" @@ -372,6 +408,42 @@ fn cvt_rn_satfinite_e4m3x2_f32_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn cvt_rn_satfinite_e5m2x2_bf16x2_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032"), + Some("v1:i1032") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::convert::cvt_rn_satfinite_e5m2x2_bf16x2" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2"), + Some("v1:i1032") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1032" + )); +} +#[test] fn cvt_rn_satfinite_e5m2x2_f16x2_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0824" @@ -444,6 +516,42 @@ fn cvt_rn_satfinite_e5m2x2_f32_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn cvt_rn_satfinite_relu_e4m3x2_bf16x2_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031"), + Some("v1:i1031") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2"), + Some("v1:i1031") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1031" + )); +} +#[test] fn cvt_rn_satfinite_relu_e4m3x2_f16x2_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0823" @@ -516,6 +624,42 @@ fn cvt_rn_satfinite_relu_e4m3x2_f32_uses_only_canonical_or_compatibility_defpath )); } #[test] +fn cvt_rn_satfinite_relu_e5m2x2_bf16x2_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033"), + Some("v1:i1033") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2"), + Some("v1:i1033") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1033" + )); +} +#[test] fn cvt_rn_satfinite_relu_e5m2x2_f16x2_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0825" 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 df75e02029..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 3963fd8c5e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 2d65aaa95e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 a228c4b86d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 d0ad70c41c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 90726a69e7..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6cf699ea72..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 7ba1abc20c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 8c2e34bd8f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 0e2925c790..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 71a2fcd9bd..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 023347d56c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated dispatch-path tests: `tma` intrinsics. @@ -529,6 +529,42 @@ fn cp_async_bulk_tensor_1d_g2s_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn cp_async_bulk_tensor_1d_g2s_cta_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034"), + Some("v1:i1034") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::tma::cp_async_bulk_tensor_1d_g2s_cta" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_1d_g2s_cta" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_1d_g2s_cta" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::tma::cp_async_bulk_tensor_1d_g2s_cta"), + Some("v1:i1034") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1034" + )); +} +#[test] fn cp_async_bulk_tensor_1d_s2g_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0335" @@ -601,6 +637,42 @@ fn cp_async_bulk_tensor_2d_g2s_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn cp_async_bulk_tensor_2d_g2s_cta_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035"), + Some("v1:i1035") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::tma::cp_async_bulk_tensor_2d_g2s_cta" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_cta" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_cta" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::tma::cp_async_bulk_tensor_2d_g2s_cta"), + Some("v1:i1035") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1035" + )); +} +#[test] fn cp_async_bulk_tensor_2d_g2s_multicast_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0330" @@ -745,6 +817,42 @@ fn cp_async_bulk_tensor_3d_g2s_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn cp_async_bulk_tensor_3d_g2s_cta_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036"), + Some("v1:i1036") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::tma::cp_async_bulk_tensor_3d_g2s_cta" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_3d_g2s_cta" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_3d_g2s_cta" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::tma::cp_async_bulk_tensor_3d_g2s_cta"), + Some("v1:i1036") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1036" + )); +} +#[test] fn cp_async_bulk_tensor_3d_s2g_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0337" @@ -817,6 +925,42 @@ fn cp_async_bulk_tensor_4d_g2s_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn cp_async_bulk_tensor_4d_g2s_cta_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037"), + Some("v1:i1037") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::tma::cp_async_bulk_tensor_4d_g2s_cta" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_4d_g2s_cta" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_4d_g2s_cta" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::tma::cp_async_bulk_tensor_4d_g2s_cta"), + Some("v1:i1037") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1037" + )); +} +#[test] fn cp_async_bulk_tensor_4d_s2g_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0338" @@ -889,6 +1033,42 @@ fn cp_async_bulk_tensor_5d_g2s_uses_only_canonical_or_compatibility_defpaths() { )); } #[test] +fn cp_async_bulk_tensor_5d_g2s_cta_uses_only_canonical_or_compatibility_defpaths() { + assert!(is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038" + )); + assert!(is_raw_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038" + )); + assert_eq!( + generated_intrinsic_marker("cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038"), + Some("v1:i1038") + ); + assert!(matches!( + classify_raw_intrinsic_path( + "cuda_intrinsics", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038".into() + ), + RawIntrinsicIdentity::Known(_) + )); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::tma::cp_async_bulk_tensor_5d_g2s_cta" + )); + assert!(is_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_5d_g2s_cta" + )); + assert!(!is_raw_generated_intrinsic_path( + "cuda_device::tma::cp_async_bulk_tensor_5d_g2s_cta" + )); + assert_eq!( + generated_intrinsic_marker("cuda_device::tma::cp_async_bulk_tensor_5d_g2s_cta"), + Some("v1:i1038") + ); + assert!(!is_generated_intrinsic_path( + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v2::i1038" + )); +} +#[test] fn cp_async_bulk_tensor_5d_s2g_uses_only_canonical_or_compatibility_defpaths() { assert!(is_generated_intrinsic_path( "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0339" 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 bf76ce6cd9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 481095a476..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 ed61420ea6..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6698b2e851..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 f587d63acd..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 7d87b7e96a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `tma` intrinsics. @@ -15,7 +15,7 @@ use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; @@ -587,6 +587,13 @@ pub(super) fn try_dispatch( "v1:i0328", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034" + | "cuda_device::tma::cp_async_bulk_tensor_1d_g2s_cta" => { + Ok(Some(super::super::tma::emit_tma_g2s_cta( + ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, 1, + "v1:i1034", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0335" | "cuda_device::tma::cp_async_bulk_tensor_1d_s2g" => { Ok(Some(super::super::tma::emit_tma_s2g( @@ -601,6 +608,13 @@ pub(super) fn try_dispatch( "v1:i0329", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035" + | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_cta" => { + Ok(Some(super::super::tma::emit_tma_g2s_cta( + ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, 2, + "v1:i1035", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0330" | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_multicast" => { Ok(Some(super::super::tma::emit_tma_g2s_multicast( @@ -627,6 +641,13 @@ pub(super) fn try_dispatch( "v1:i0332", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036" + | "cuda_device::tma::cp_async_bulk_tensor_3d_g2s_cta" => { + Ok(Some(super::super::tma::emit_tma_g2s_cta( + ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, 3, + "v1:i1036", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0337" | "cuda_device::tma::cp_async_bulk_tensor_3d_s2g" => { Ok(Some(super::super::tma::emit_tma_s2g( @@ -641,6 +662,13 @@ pub(super) fn try_dispatch( "v1:i0333", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037" + | "cuda_device::tma::cp_async_bulk_tensor_4d_g2s_cta" => { + Ok(Some(super::super::tma::emit_tma_g2s_cta( + ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, 4, + "v1:i1037", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0338" | "cuda_device::tma::cp_async_bulk_tensor_4d_s2g" => { Ok(Some(super::super::tma::emit_tma_s2g( @@ -655,6 +683,13 @@ pub(super) fn try_dispatch( "v1:i0334", )?)) } + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038" + | "cuda_device::tma::cp_async_bulk_tensor_5d_g2s_cta" => { + Ok(Some(super::super::tma::emit_tma_g2s_cta( + ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, 5, + "v1:i1038", + )?)) + } "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0339" | "cuda_device::tma::cp_async_bulk_tensor_5d_s2g" => { Ok(Some(super::super::tma::emit_tma_s2g( 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 016a3ae61b..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated raw/compatibility dispatch arms: `vote` intrinsics. @@ -15,7 +15,7 @@ use crate::translator::{rvalue, terminator::helpers, values::ValueMap}; use dialect_nvvm::ops::{VoteSyncAllOp, VoteSyncAnyOp, VoteSyncBallotOp, VoteSyncUniOp}; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 ea2efec42c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -16,7 +16,7 @@ use dialect_nvvm::ops::BarWarpSyncOp; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 7da0eb23dd..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -17,7 +17,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 f445adee0d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -19,7 +19,7 @@ use dialect_nvvm::ops::{ }; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; 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 966b12094d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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. @@ -19,7 +19,7 @@ use dialect_nvvm::ops::{ use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::operation::Operation; use rustc_public::mir; diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/iket.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/iket.rs index 78d0bcecbd..45d28bc4fa 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/iket.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/iket.rs @@ -16,7 +16,7 @@ use pliron::{ basic_block::BasicBlock, context::{Context, Ptr}, input_err, - location::{Located, Location}, + location::Location, op::Op, operation::Operation, value::Value, diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/indexing.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/indexing.rs index 6ca3b5e837..bd175032f3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/indexing.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/indexing.rs @@ -39,7 +39,7 @@ use crate::translator::values::ValueMap; use pliron::basic_block::BasicBlock; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::Typed; diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/integer_minmax.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/integer_minmax.rs new file mode 100644 index 0000000000..f02c3da649 --- /dev/null +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/integer_minmax.rs @@ -0,0 +1,370 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Primitive integer `Ord::min` and `Ord::max` intrinsics. + +use super::super::helpers; +use crate::error::{TranslationErr, TranslationResult}; +use crate::translator::{rvalue, types, values::ValueMap}; +use dialect_mir::{ + ops::{MirCallOp, MirLtOp}, + rust_intrinsics, +}; +use pliron::basic_block::BasicBlock; +use pliron::builtin::{ + attributes::StringAttr, + 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; + +/// Integer ordering intrinsics emitted by libcore. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RustIntegerMinMaxIntrinsic { + /// Select the smaller operand. + Min, + /// Select the larger operand. + Max, +} + +impl RustIntegerMinMaxIntrinsic { + /// Recognize exact libcore intrinsic paths, including the std reexports. + pub fn from_core_path(name: &str) -> Option { + match name { + "core::intrinsics::integer_min" | "std::intrinsics::integer_min" => Some(Self::Min), + "core::intrinsics::integer_max" | "std::intrinsics::integer_max" => Some(Self::Max), + _ => None, + } + } +} + +/// Match core's specified fallback: `a < b` selects `(a, b)` for min and +/// `(b, a)` for max. Keeping the original integer types lets existing compare +/// lowering recover signedness, without narrowing any width (including 128). +fn emit_integer_minmax( + ctx: &mut Context, + intrinsic: RustIntegerMinMaxIntrinsic, + args: &[Value], + return_type: TypeHandle, + block: Ptr, + prev_op: Option>, + loc: Location, +) -> TranslationResult> { + let [lhs, rhs] = args else { + return input_err!( + loc, + TranslationErr::type_error("integer min/max expects two operands") + ); + }; + let valid_integer = return_type + .deref(ctx) + .downcast_ref::() + .is_some_and(|ty| { + matches!(ty.width(), 8 | 16 | 32 | 64 | 128) + && matches!(ty.signedness(), Signedness::Signed | Signedness::Unsigned) + }); + if !valid_integer || lhs.get_type(ctx) != return_type || rhs.get_type(ctx) != return_type { + return input_err!( + loc, + TranslationErr::type_error( + "integer min/max requires matching primitive integer operands and result" + ) + ); + } + + let bool_type = types::get_bool_type(ctx).into(); + let cmp = Operation::new( + ctx, + MirLtOp::get_concrete_op_info(), + vec![bool_type], + vec![*lhs, *rhs], + vec![], + 0, + ); + cmp.deref_mut(ctx).set_loc(loc.clone()); + helpers::insert_op(ctx, cmp, block, prev_op); + let condition = cmp.deref(ctx).get_result(0); + let (if_true, if_false) = match intrinsic { + RustIntegerMinMaxIntrinsic::Min => (*lhs, *rhs), + RustIntegerMinMaxIntrinsic::Max => (*rhs, *lhs), + }; + // This existing placeholder lowers to an ordinary LLVM select; it carries + // no unpredictable metadata or other optimization promises. + let select = Operation::new( + ctx, + MirCallOp::get_concrete_op_info(), + vec![return_type], + vec![condition, if_true, if_false], + vec![], + 0, + ); + select.deref_mut(ctx).set_loc(loc); + select.deref_mut(ctx).attributes.set( + "callee".try_into().unwrap(), + StringAttr::new(rust_intrinsics::CALLEE_SELECT_UNPREDICTABLE.into()), + ); + select.insert_after(ctx, cmp); + Ok(select) +} + +/// Translate integer min/max using the ordinary compare and select lowering. +pub fn emit_rust_integer_minmax_intrinsic( + ctx: &mut Context, + body: &mir::Body, + intrinsic: RustIntegerMinMaxIntrinsic, + args: &[mir::Operand], + destination: &mir::Place, + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, +) -> TranslationResult> { + let return_type = types::translate_destination_type(ctx, body, destination, &loc)?; + let mut values = Vec::with_capacity(args.len()); + let mut last_op = prev_op; + for arg in args { + let (value, last) = + rvalue::translate_operand(ctx, body, arg, value_map, block_ptr, last_op, loc.clone())?; + values.push(value); + last_op = last; + } + let (prepared, last_op) = helpers::prepare_destination_write( + ctx, + body, + destination, + value_map, + block_ptr, + last_op, + loc.clone(), + )?; + let select = emit_integer_minmax( + ctx, + intrinsic, + &values, + return_type, + block_ptr, + last_op, + loc.clone(), + )?; + let result = select.deref(ctx).get_result(0); + helpers::emit_prepared_result_and_goto( + ctx, + prepared, + result, + target, + block_ptr, + select, + value_map, + block_map, + loc, + "integer min/max call without target block", + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use llvm_export::{attributes::ICmpPredicateAttr, ops as llvm}; + use pliron::builtin::{ + attributes::TypeAttr, op_interfaces::SymbolOpInterface, ops::ModuleOp, types::FunctionType, + }; + use pliron::common_traits::Verify; + use pliron::linked_list::ContainsLinkedList; + + #[test] + fn recognizes_only_integer_intrinsic_paths() { + for prefix in ["core", "std"] { + assert_eq!( + RustIntegerMinMaxIntrinsic::from_core_path(&format!( + "{prefix}::intrinsics::integer_min" + )), + Some(RustIntegerMinMaxIntrinsic::Min) + ); + assert_eq!( + RustIntegerMinMaxIntrinsic::from_core_path(&format!( + "{prefix}::intrinsics::integer_max" + )), + Some(RustIntegerMinMaxIntrinsic::Max) + ); + } + for path in [ + "user::intrinsics::integer_min", + "core::intrinsics::integer_minimum", + "core::intrinsics::integer_max_extra", + "core::cmp::min", + "core::intrinsics::minimum_number_nsz", + ] { + assert_eq!(RustIntegerMinMaxIntrinsic::from_core_path(path), None); + } + } + + #[test] + fn lowers_minmax_with_original_width_signedness_and_selected_operands() { + for width in [8, 16, 32, 64, 128] { + for (signedness, predicate) in [ + (Signedness::Signed, ICmpPredicateAttr::SLT), + (Signedness::Unsigned, ICmpPredicateAttr::ULT), + ] { + for (intrinsic, true_index, false_index) in [ + (RustIntegerMinMaxIntrinsic::Min, 0, 1), + (RustIntegerMinMaxIntrinsic::Max, 1, 0), + ] { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + mir_lower::register(&mut ctx); + let ty = IntegerType::get(&ctx, width, signedness).into(); + let module = + ModuleOp::new(&mut ctx, "minmax".try_into().unwrap()).get_operation(); + let func_type = FunctionType::get(&ctx, vec![ty, ty], vec![ty]); + let func = Operation::new( + &mut ctx, + dialect_mir::ops::MirFuncOp::get_concrete_op_info(), + vec![], + vec![], + vec![], + 1, + ); + let mir_func = dialect_mir::ops::MirFuncOp::new( + &mut ctx, + func, + TypeAttr::new(func_type.into()), + ); + mir_func.set_symbol_name(&mut ctx, "minmax".try_into().unwrap()); + let module_region = module.deref(&ctx).get_region(0); + let module_block = module_region.deref(&ctx).iter(&ctx).next().unwrap(); + func.insert_at_back(module_block, &ctx); + let region = func.deref(&ctx).get_region(0); + let block = BasicBlock::new(&mut ctx, None, vec![ty, ty]); + block.insert_at_back(region, &ctx); + let args = [ + block.deref(&ctx).get_argument(0), + block.deref(&ctx).get_argument(1), + ]; + let select = emit_integer_minmax( + &mut ctx, + intrinsic, + &args, + ty, + block, + None, + Location::Unknown, + ) + .unwrap(); + let result = select.deref(&ctx).get_result(0); + let ret = Operation::new( + &mut ctx, + dialect_mir::ops::MirReturnOp::get_concrete_op_info(), + vec![], + vec![result], + vec![], + 0, + ); + ret.insert_after(&ctx, select); + mir_lower::lower_mir_to_llvm(&mut ctx, module).unwrap(); + module.deref(&ctx).verify(&ctx).unwrap(); + let func = module_block.deref(&ctx).iter(&ctx).next().unwrap(); + let region = func.deref(&ctx).get_region(0); + let operations: Vec<_> = region + .deref(&ctx) + .iter(&ctx) + .flat_map(|block| block.deref(&ctx).iter(&ctx)) + .collect(); + let compare = operations + .iter() + .find_map(|op| Operation::get_op::(*op, &ctx)) + .unwrap(); + assert_eq!(compare.predicate(&ctx), predicate); + let compare = compare.get_operation(); + let select = operations + .iter() + .find_map(|op| Operation::get_op::(*op, &ctx)) + .unwrap() + .get_operation(); + assert_eq!( + select.deref(&ctx).get_operand(0), + compare.deref(&ctx).get_result(0) + ); + assert_eq!( + select.deref(&ctx).get_operand(1), + block.deref(&ctx).get_argument(true_index) + ); + assert_eq!( + select.deref(&ctx).get_operand(2), + block.deref(&ctx).get_argument(false_index) + ); + assert_eq!( + compare.deref(&ctx).get_operand(0), + block.deref(&ctx).get_argument(0) + ); + assert_eq!( + compare.deref(&ctx).get_operand(1), + block.deref(&ctx).get_argument(1) + ); + assert_eq!( + select + .deref(&ctx) + .get_result(0) + .get_type(&ctx) + .deref(&ctx) + .downcast_ref::() + .unwrap() + .width(), + width + ); + } + } + } + } + + #[test] + fn rejects_invalid_integer_minmax_signatures_before_emission() { + let mut ctx = Context::new(); + dialect_mir::register(&mut ctx); + let integer: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let signed: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signed).into(); + let wide: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let boolean: TypeHandle = IntegerType::get(&ctx, 1, Signedness::Signless).into(); + let signless: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let odd_width: TypeHandle = IntegerType::get(&ctx, 24, Signedness::Unsigned).into(); + let float: TypeHandle = pliron::builtin::types::FP32Type::get(&ctx).into(); + for (args, result) in [ + (vec![integer], integer), + (vec![integer, integer, integer], integer), + (vec![integer, wide], integer), + (vec![integer, signed], integer), + (vec![integer, integer], wide), + (vec![boolean, boolean], boolean), + (vec![signless, signless], signless), + (vec![odd_width, odd_width], odd_width), + (vec![float, float], float), + ] { + let block = BasicBlock::new(&mut ctx, None, args.clone()); + let values: Vec<_> = (0..args.len()) + .map(|index| block.deref(&ctx).get_argument(index)) + .collect(); + assert!( + emit_integer_minmax( + &mut ctx, + RustIntegerMinMaxIntrinsic::Min, + &values, + result, + block, + None, + Location::Unknown + ) + .is_err() + ); + assert!(block.deref(&ctx).iter(&ctx).next().is_none()); + } + } +} diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/layout.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/layout.rs index cf9216d1a1..bbed504db2 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/layout.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/layout.rs @@ -31,7 +31,7 @@ use dialect_mir::types::MirSliceType; use pliron::basic_block::BasicBlock; use pliron::builtin::attributes::IntegerAttr; use pliron::context::{Context, Ptr}; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::Typed; 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 2435f05f34..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 @@ -19,7 +19,7 @@ use pliron::builtin::types::IntegerType; use pliron::common_traits::Verify; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::Typed; @@ -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/intrinsics/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/mod.rs index f70f591f38..5c50b07cde 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/mod.rs @@ -12,6 +12,7 @@ //! |--------------|------------------------------------------------------------------------------| //! | `generated` | Every catalog intrinsic, dispatched by canonical and compatibility path | //! | `indexing` | `threadIdx_*`, `blockIdx_*`, `index_1d`, `index_2d::`, `index_2d_runtime` | +//! | `integer_minmax` | Rust compiler integer minimum and maximum intrinsics | //! | `memory` | `SharedArray`, `stmatrix_*`, type conversions | //! | `atomic` | Atomic read-modify-write and compare-exchange | //! | `wgmma` | Hopper WGMMA matrix operations | @@ -51,6 +52,7 @@ pub mod float_math; pub mod generated; pub mod iket; pub mod indexing; +pub mod integer_minmax; pub mod layout; pub mod memory; pub mod saturating; diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/tma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/tma.rs index 054f33755c..034567db76 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/tma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/tma.rs @@ -13,7 +13,9 @@ use crate::translator::rvalue; use crate::translator::values::ValueMap; use dialect_mir::ops::MirConstantOp; use dialect_nvvm::ops::{ - CpAsyncBulkCommitGroupOp, CpAsyncBulkTensorG2sTile1dOp, + CpAsyncBulkCommitGroupOp, CpAsyncBulkTensorG2sCtaTile1dOp, CpAsyncBulkTensorG2sCtaTile2dOp, + CpAsyncBulkTensorG2sCtaTile3dOp, CpAsyncBulkTensorG2sCtaTile4dOp, + CpAsyncBulkTensorG2sCtaTile5dOp, CpAsyncBulkTensorG2sTile1dOp, CpAsyncBulkTensorG2sTile2dMulticastCg2Op, CpAsyncBulkTensorG2sTile2dMulticastOp, CpAsyncBulkTensorG2sTile2dOp, CpAsyncBulkTensorG2sTile3dOp, CpAsyncBulkTensorG2sTile4dOp, CpAsyncBulkTensorG2sTile5dOp, CpAsyncBulkTensorS2gTile1dOp, CpAsyncBulkTensorS2gTile2dOp, @@ -24,7 +26,7 @@ use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use rustc_public::mir; @@ -61,6 +63,44 @@ pub fn emit_tma_g2s( loc: Location, dims: usize, marker: &str, +) -> TranslationResult> { + emit_tma_g2s_impl( + ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, dims, marker, false, + ) +} + +/// Emits a TMA G2S form whose destination belongs to the issuing CTA. +pub fn emit_tma_g2s_cta( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + dims: usize, + marker: &str, +) -> TranslationResult> { + emit_tma_g2s_impl( + ctx, body, args, target, block_ptr, prev_op, value_map, block_map, loc, dims, marker, true, + ) +} + +fn emit_tma_g2s_impl( + ctx: &mut Context, + body: &mir::Body, + args: &[mir::Operand], + target: &Option, + block_ptr: Ptr, + prev_op: Option>, + value_map: &mut ValueMap, + block_map: &[Ptr], + loc: Location, + dims: usize, + marker: &str, + cta: bool, ) -> TranslationResult> { // Expected args: dst, tensor_map, coord0, [coord1, ...], barrier let expected_args = 3 + dims; // dst + tensor_map + coords + barrier @@ -68,8 +108,9 @@ pub fn emit_tma_g2s( return input_err!( loc.clone(), TranslationErr::unsupported(format!( - "cp_async_bulk_tensor_{}d_g2s expects {} arguments, got {}", + "cp_async_bulk_tensor_{}d_g2s{} expects {} arguments, got {}", dims, + if cta { "_cta" } else { "" }, expected_args, args.len() )) @@ -135,38 +176,38 @@ pub fn emit_tma_g2s( last_op = last_op_after; } - // Add default cta_mask (i16 = 0) and cache_hint (i64 = 0) + // Cluster G2S carries a default CTA mask. CTA-local G2S has no mask. + // Both forms carry the default cache hint. use pliron::utils::apint::APInt; use std::num::NonZeroUsize; - let i16_type = IntegerType::get(ctx, 16, Signedness::Signed); let i64_type = IntegerType::get(ctx, 64, Signedness::Unsigned); - // Create constant for cta_mask = 0 - let cta_mask_apint = APInt::from_i64(0, NonZeroUsize::new(16).unwrap()); - let cta_mask_attr = pliron::builtin::attributes::IntegerAttr::new(i16_type, cta_mask_apint); - - let cta_mask_raw_op = Operation::new( - ctx, - MirConstantOp::get_concrete_op_info(), - vec![i16_type.to_handle()], - vec![], - vec![], - 0, - ); - cta_mask_raw_op.deref_mut(ctx).set_loc(loc.clone()); - let cta_mask_const = MirConstantOp::new(cta_mask_raw_op); - cta_mask_const.set_attr_value(ctx, cta_mask_attr); - - if let Some(prev) = last_op { - cta_mask_const.get_operation().insert_after(ctx, prev); - } else { - cta_mask_const - .get_operation() - .insert_at_front(block_ptr, ctx); + if !cta { + let i16_type = IntegerType::get(ctx, 16, Signedness::Signed); + let cta_mask_apint = APInt::from_i64(0, NonZeroUsize::new(16).unwrap()); + let cta_mask_attr = pliron::builtin::attributes::IntegerAttr::new(i16_type, cta_mask_apint); + let cta_mask_raw_op = Operation::new( + ctx, + MirConstantOp::get_concrete_op_info(), + vec![i16_type.to_handle()], + vec![], + vec![], + 0, + ); + cta_mask_raw_op.deref_mut(ctx).set_loc(loc.clone()); + let cta_mask_const = MirConstantOp::new(cta_mask_raw_op); + cta_mask_const.set_attr_value(ctx, cta_mask_attr); + if let Some(prev) = last_op { + cta_mask_const.get_operation().insert_after(ctx, prev); + } else { + cta_mask_const + .get_operation() + .insert_at_front(block_ptr, ctx); + } + operands.push(cta_mask_const.get_operation().deref(ctx).get_result(0)); + last_op = Some(cta_mask_const.get_operation()); } - let cta_mask = cta_mask_const.get_operation().deref(ctx).get_result(0); - operands.push(cta_mask); // Create constant for cache_hint = 0 let cache_hint_apint = APInt::from_i64(0, NonZeroUsize::new(64).unwrap()); @@ -183,19 +224,28 @@ pub fn emit_tma_g2s( cache_hint_raw_op.deref_mut(ctx).set_loc(loc.clone()); let cache_hint_const = MirConstantOp::new(cache_hint_raw_op); cache_hint_const.set_attr_value(ctx, cache_hint_attr); - cache_hint_const - .get_operation() - .insert_after(ctx, cta_mask_const.get_operation()); + if let Some(prev) = last_op { + cache_hint_const.get_operation().insert_after(ctx, prev); + } else { + cache_hint_const + .get_operation() + .insert_at_front(block_ptr, ctx); + } let cache_hint = cache_hint_const.get_operation().deref(ctx).get_result(0); operands.push(cache_hint); // Select the appropriate NVVM op based on dimensions let op_id = match dims { + 1 if cta => CpAsyncBulkTensorG2sCtaTile1dOp::get_concrete_op_info(), 1 => CpAsyncBulkTensorG2sTile1dOp::get_concrete_op_info(), + 2 if cta => CpAsyncBulkTensorG2sCtaTile2dOp::get_concrete_op_info(), 2 => CpAsyncBulkTensorG2sTile2dOp::get_concrete_op_info(), + 3 if cta => CpAsyncBulkTensorG2sCtaTile3dOp::get_concrete_op_info(), 3 => CpAsyncBulkTensorG2sTile3dOp::get_concrete_op_info(), + 4 if cta => CpAsyncBulkTensorG2sCtaTile4dOp::get_concrete_op_info(), 4 => CpAsyncBulkTensorG2sTile4dOp::get_concrete_op_info(), + 5 if cta => CpAsyncBulkTensorG2sCtaTile5dOp::get_concrete_op_info(), 5 => CpAsyncBulkTensorG2sTile5dOp::get_concrete_op_info(), _ => { return input_err!( @@ -213,7 +263,7 @@ pub fn emit_tma_g2s( ctx, op_id, vec![], // No results - operands, // dst, barrier, tensor_map, coords..., cta_mask, cache_hint + operands, // dst, barrier, tensor_map, coords..., optional cta_mask, cache_hint vec![], 0, ); diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/wgmma.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/wgmma.rs index 4d662072a4..62b0ccb3cb 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/wgmma.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/intrinsics/wgmma.rs @@ -19,7 +19,7 @@ use pliron::basic_block::BasicBlock; use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; use pliron::input_err; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::value::Value; 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 34f8d9889c..524771dea8 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs @@ -77,7 +77,7 @@ use pliron::builtin::types::{FunctionType, IntegerType, Signedness}; use pliron::context::{Context, Ptr}; use pliron::identifier::Legaliser; use pliron::linked_list::ContainsLinkedList; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; @@ -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. @@ -1254,25 +1273,47 @@ fn translate_call( return Ok(result); } - // Handle diverging calls (calls that never return, like unwrap_failed, panic, etc.) - // These have no target block because the function never returns. + // `Call { target: None }` means there is no normal successor. rustc + // clears the target when the destination type is visibly uninhabited: + // `!`, `core::convert::Infallible`, an empty enum. That does not mean the + // callee can be erased. An emitted Rust function with such a return type + // may do observable work before it diverges. When the resolved instance + // is in this pipeline run's emitted-function set, keep the call and let + // `emit_function_call` end the block with `mir.unreachable`, which is + // what rustc's own backend emits for a call without a destination. // - // The callee is dropped and replaced by an immediate trap. That applies - // to every `-> !` callee reaching this point, including a user - // `#[device]` fn that legitimately diverges (e.g. `loop {}`); full Rust - // fidelity would emit the call for resolvable callees the way - // `translate_function_item_call` does. Dropping the call is a - // pre-existing semantic: the trap only makes it safe, where the bare - // `unreachable` previously emitted here was UB that let `opt` delete - // the whole panic path. - // - // Panic entry points additionally never get here with any statements - // translated ahead of them: `block::translate_block` recognizes the same - // shape via [`is_dropped_panic_call`] and emits the trap directly. + // Every other diverging call keeps the pre-existing trap fallback. That + // covers callees the collector skipped, such as core helpers without MIR + // and local `-> !` bodies that are nothing but a panic, and a call + // through a function pointer, whose `func` operand is not a constant and + // so resolves no instance. Panic entry points never reach this arm: for + // a callee whose path matches `is_panic_entry_path`, + // `block::translate_block` lowers the whole block to the trap through + // `is_dropped_panic_call` before any call is classified. Keep this arm + // before the unsupported-intrinsic and libm checks below so the trapped + // cases take the same path as before. if target_usize.is_none() { - // This is a diverging call (returns !) - emit trap + unreachable - // Examples: unwrap_failed(), panic!(), abort() - return Ok(emit_trap_unreachable_after(ctx, block_ptr, prev_op, loc)); + let resolved_instance = match func { + mir::Operand::Constant(constant) + if *constant.const_.kind() == ConstantKind::ZeroSized => + { + match constant.const_.ty().kind() { + rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::FnDef( + fn_def, + substs, + )) => rustc_public::mir::mono::Instance::resolve(fn_def, &substs).ok(), + _ => None, + } + } + _ => None, + }; + + if !resolved_instance + .as_ref() + .is_some_and(crate::translator::facts::is_emitted_function) + { + return Ok(emit_trap_unreachable_after(ctx, block_ptr, prev_op, loc)); + } } // A call to a rustc intrinsic that no dispatch arm above recognized can @@ -2089,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, })) } @@ -2761,6 +2803,26 @@ fn try_dispatch_intrinsic( )?)); } + if let Some(intrinsic) = + intrinsics::integer_minmax::RustIntegerMinMaxIntrinsic::from_core_path(name) + { + return Ok(Some( + intrinsics::integer_minmax::emit_rust_integer_minmax_intrinsic( + ctx, + body, + intrinsic, + args, + destination, + target, + block_ptr, + prev_op, + value_map, + block_map, + loc, + )?, + )); + } + if let Some(intrinsic) = intrinsics::saturating::RustSaturatingIntrinsic::from_core_path(name) { return Ok(Some( intrinsics::saturating::emit_rust_saturating_intrinsic( @@ -2827,7 +2889,13 @@ fn try_dispatch_intrinsic( )?)); } - if let Some(intrinsic) = intrinsics::float_math::RustFloatMathIntrinsic::from_core_path(name) { + if let Some(intrinsic) = + intrinsics::float_math::RustFloatMathIntrinsic::from_core_path_with_substs( + name, + type_substs, + ) + .map_err(|error| input_error!(loc.clone(), TranslationErr::unsupported(error)))? + { return Ok(Some( intrinsics::float_math::emit_rust_float_math_intrinsic( ctx, @@ -3294,6 +3362,7 @@ fn try_dispatch_intrinsic( block_map, loc, 64, + false, )?, )), "cuda_device::shared::cvta_generic_to_shared_u32" => Ok(Some( @@ -3309,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 ee5c23b8e7..e3ad890548 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/types.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/types.rs @@ -364,6 +364,14 @@ fn translate_pointer_like( let elem = shared_array_element_type(ctx, &substs, "SharedArray")?; Ok(facts::mint_shared_ptr_type(ctx, elem, origin).into()) } + rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Adt(adt_def, substs)) + if is_cuda_device_adt(&adt_def, "SharedSlot") => + { + // `*mut SharedSlot` is the field of `SharedPtr`: a pointer + // to a `T` in CTA-shared memory, by its type. + let elem = shared_array_element_type(ctx, &substs, "SharedSlot")?; + Ok(facts::mint_shared_ptr_type(ctx, elem, origin).into()) + } rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Adt(adt_def, _substs)) if is_cuda_device_adt(&adt_def, "Barrier") => { @@ -662,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(); @@ -718,6 +749,14 @@ pub fn translate_type( // ThreadIndex is a newtype around usize - translate to usize // The type safety is enforced at the Rust level, not the IR level Ok(get_usize_type(ctx).into()) + } else if is_cuda_device_adt(&adt_def, "SharedPtr") { + // `SharedPtr { ptr: *mut SharedSlot }` is a scalar-lowered + // newtype: the value is its field, a CTA-shared (`addrspace(3)`) + // pointer to `T`. Params, returns, locals and fields of type + // `SharedPtr` are therefore shared pointers by type, across + // function boundaries. + let elem = shared_array_element_type(ctx, &substs, "SharedPtr")?; + Ok(facts::mint_shared_ptr_type(ctx, elem, facts::abi_shared_ptr_field()).into()) } else if is_cuda_device_adt(&adt_def, "RangeToken") { Ok(dialect_iket::types::IketRangeTokenType::get(ctx).into()) } else if is_cuda_device_adt(&adt_def, "SharedArray") { @@ -748,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(); 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/context.rs b/cuda-oxide/crates/mir-lower/src/context.rs index f4fa8112ce..88c6f73388 100644 --- a/cuda-oxide/crates/mir-lower/src/context.rs +++ b/cuda-oxide/crates/mir-lower/src/context.rs @@ -94,7 +94,7 @@ pub struct SharedGlobalDeclaration { pub mir_elem_type: TypeHandle, /// Number of elements in the shared allocation. pub size: u64, - /// Explicit alignment, or zero when natural alignment is requested. + /// Effective physical storage alignment in bytes. pub alignment: u64, /// Source-level debug identity carried on the allocation, when present. pub debug_info: Option, 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 95b3f0daf5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fc7b13775f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 79e59a186a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `cluster_barrier` intrinsics. @@ -18,7 +18,6 @@ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 809bdd2fe6..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 30ec31ca6c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 cc37b0bcac..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `debug_control` intrinsics. @@ -17,7 +17,6 @@ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 5f1c71aa9e..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 49d0485ba0..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 eb33f917ed..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 64976d1c34..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `extended_minmax` intrinsics. @@ -20,7 +20,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 dbfe47cc8c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 3c029721fe..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `ldmatrix` intrinsics. @@ -19,7 +19,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 760b526608..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6a2e00d2da..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fdaa4bf21c..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 ce5d980b17..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 3d9d39e5ca..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 8b10729fce..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `packed_atomic` intrinsics. @@ -19,7 +19,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 7f9e1b49fb..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `packed_conversion` intrinsics. @@ -16,8 +16,10 @@ use crate::convert::intrinsics::{ use dialect_nvvm::ops::{ CvtF16x2F32Op, CvtF32x2Bf16x2Op, CvtRnF16x2E4m3x2Op, CvtRnF16x2E5m2x2Op, CvtRnReluBf16x2F32Op, CvtRnReluF16x2E4m3x2Op, CvtRnReluF16x2E5m2x2Op, CvtRnReluF16x2F32Op, - CvtRnSatfiniteE4m3x2F16x2Op, CvtRnSatfiniteE4m3x2F32Op, CvtRnSatfiniteE5m2x2F16x2Op, - CvtRnSatfiniteE5m2x2F32Op, CvtRnSatfiniteReluE4m3x2F16x2Op, CvtRnSatfiniteReluE4m3x2F32Op, + CvtRnSatfiniteE4m3x2Bf16x2Op, CvtRnSatfiniteE4m3x2F16x2Op, CvtRnSatfiniteE4m3x2F32Op, + CvtRnSatfiniteE5m2x2Bf16x2Op, CvtRnSatfiniteE5m2x2F16x2Op, CvtRnSatfiniteE5m2x2F32Op, + CvtRnSatfiniteReluE4m3x2Bf16x2Op, CvtRnSatfiniteReluE4m3x2F16x2Op, + CvtRnSatfiniteReluE4m3x2F32Op, CvtRnSatfiniteReluE5m2x2Bf16x2Op, CvtRnSatfiniteReluE5m2x2F16x2Op, CvtRnSatfiniteReluE5m2x2F32Op, CvtRzBf16x2F32Op, CvtRzF16x2F32Op, }; @@ -179,6 +181,25 @@ impl MirToLlvmConversion for CvtRnReluF16x2F32Op { } } +#[op_interface_impl] +impl MirToLlvmConversion for CvtRnSatfiniteE4m3x2Bf16x2Op { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + _operands_info: &OperandsInfo, + ) -> Result<()> { + convert_generated_packed_unary( + ctx, + rewriter, + self.get_operation(), + "cvt.rn.satfinite.e4m3x2.bf16x2", + 16, + 32, + ) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CvtRnSatfiniteE4m3x2F16x2Op { fn convert( @@ -217,6 +238,25 @@ impl MirToLlvmConversion for CvtRnSatfiniteE4m3x2F32Op { } } +#[op_interface_impl] +impl MirToLlvmConversion for CvtRnSatfiniteE5m2x2Bf16x2Op { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + _operands_info: &OperandsInfo, + ) -> Result<()> { + convert_generated_packed_unary( + ctx, + rewriter, + self.get_operation(), + "cvt.rn.satfinite.e5m2x2.bf16x2", + 16, + 32, + ) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CvtRnSatfiniteE5m2x2F16x2Op { fn convert( @@ -255,6 +295,25 @@ impl MirToLlvmConversion for CvtRnSatfiniteE5m2x2F32Op { } } +#[op_interface_impl] +impl MirToLlvmConversion for CvtRnSatfiniteReluE4m3x2Bf16x2Op { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + _operands_info: &OperandsInfo, + ) -> Result<()> { + convert_generated_packed_unary( + ctx, + rewriter, + self.get_operation(), + "cvt.rn.satfinite.relu.e4m3x2.bf16x2", + 16, + 32, + ) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CvtRnSatfiniteReluE4m3x2F16x2Op { fn convert( @@ -293,6 +352,25 @@ impl MirToLlvmConversion for CvtRnSatfiniteReluE4m3x2F32Op { } } +#[op_interface_impl] +impl MirToLlvmConversion for CvtRnSatfiniteReluE5m2x2Bf16x2Op { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + _operands_info: &OperandsInfo, + ) -> Result<()> { + convert_generated_packed_unary( + ctx, + rewriter, + self.get_operation(), + "cvt.rn.satfinite.relu.e5m2x2.bf16x2", + 16, + 32, + ) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CvtRnSatfiniteReluE5m2x2F16x2Op { fn convert( 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 7fcacbd0b9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `prmt` intrinsics. @@ -15,7 +15,6 @@ use dialect_nvvm::ops::{PrmtModeAttr, PrmtOp}; use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 a776f60f9d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 3258cf9bd4..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `register_mma` intrinsics. @@ -22,7 +22,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; @@ -2138,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 eb2f26f21f..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_arithmetic` intrinsics. @@ -18,7 +18,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 0d5779dab4..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_conversion` intrinsics. @@ -17,7 +17,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 329b2242f7..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `scalar_math` intrinsics. @@ -18,7 +18,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 6d6b104715..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `sparse_mma` intrinsics. @@ -20,7 +20,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 e8c341fbba..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 639c29ce43..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 0385565067..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6819daab9d..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `tcgen05` intrinsics. @@ -89,7 +89,6 @@ use dialect_nvvm::ops::{ use pliron::context::Context; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; -use pliron::location::Located; use pliron::op::Op; use pliron::result::Result; 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 5ba6cbc944..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,17 +4,17 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `tma` intrinsics. use crate::conversion_interface::MirToLlvmConversion; use crate::convert::intrinsics::{ tma::PrefetchTileConfig, tma::ReduceConfig, tma::convert_control, tma::convert_g2s, - tma::convert_g2s_multicast_cg2, tma::convert_prefetch_tensormap, tma::convert_prefetch_tile, - tma::convert_reduce_s2g, tma::convert_s2g, tma::convert_tensormap_fence, - tma::convert_tensormap_replace, + tma::convert_g2s_cta, tma::convert_g2s_multicast_cg2, tma::convert_prefetch_tensormap, + tma::convert_prefetch_tile, tma::convert_reduce_s2g, tma::convert_s2g, + tma::convert_tensormap_fence, tma::convert_tensormap_replace, }; use dialect_nvvm::ops::{ CpAsyncBulkCommitGroupOp, CpAsyncBulkPrefetchTensor1dL2CacheHintOp, @@ -23,7 +23,9 @@ use dialect_nvvm::ops::{ CpAsyncBulkPrefetchTensor3dL2Op, CpAsyncBulkPrefetchTensor4dL2CacheHintOp, CpAsyncBulkPrefetchTensor4dL2Op, CpAsyncBulkPrefetchTensor5dL2CacheHintOp, CpAsyncBulkPrefetchTensor5dL2Op, CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2CacheHintOp, - CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2Op, CpAsyncBulkTensorG2sTile1dOp, + CpAsyncBulkPrefetchTensorGather4TwoDimensionalL2Op, CpAsyncBulkTensorG2sCtaTile1dOp, + CpAsyncBulkTensorG2sCtaTile2dOp, CpAsyncBulkTensorG2sCtaTile3dOp, + CpAsyncBulkTensorG2sCtaTile4dOp, CpAsyncBulkTensorG2sCtaTile5dOp, CpAsyncBulkTensorG2sTile1dOp, CpAsyncBulkTensorG2sTile2dMulticastCg2Op, CpAsyncBulkTensorG2sTile2dMulticastOp, CpAsyncBulkTensorG2sTile2dOp, CpAsyncBulkTensorG2sTile3dOp, CpAsyncBulkTensorG2sTile4dOp, CpAsyncBulkTensorG2sTile5dOp, CpAsyncBulkTensorReduceAddIm2col3dOp, @@ -383,6 +385,18 @@ impl MirToLlvmConversion for CpAsyncBulkTensorG2sTile1dOp { } } +#[op_interface_impl] +impl MirToLlvmConversion for CpAsyncBulkTensorG2sCtaTile1dOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + operands_info: &OperandsInfo, + ) -> Result<()> { + convert_g2s_cta(ctx, rewriter, self.get_operation(), operands_info, 1) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CpAsyncBulkTensorS2gTile1dOp { fn convert( @@ -407,6 +421,18 @@ impl MirToLlvmConversion for CpAsyncBulkTensorG2sTile2dOp { } } +#[op_interface_impl] +impl MirToLlvmConversion for CpAsyncBulkTensorG2sCtaTile2dOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + operands_info: &OperandsInfo, + ) -> Result<()> { + convert_g2s_cta(ctx, rewriter, self.get_operation(), operands_info, 2) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CpAsyncBulkTensorG2sTile2dMulticastOp { fn convert( @@ -455,6 +481,18 @@ impl MirToLlvmConversion for CpAsyncBulkTensorG2sTile3dOp { } } +#[op_interface_impl] +impl MirToLlvmConversion for CpAsyncBulkTensorG2sCtaTile3dOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + operands_info: &OperandsInfo, + ) -> Result<()> { + convert_g2s_cta(ctx, rewriter, self.get_operation(), operands_info, 3) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CpAsyncBulkTensorS2gTile3dOp { fn convert( @@ -479,6 +517,18 @@ impl MirToLlvmConversion for CpAsyncBulkTensorG2sTile4dOp { } } +#[op_interface_impl] +impl MirToLlvmConversion for CpAsyncBulkTensorG2sCtaTile4dOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + operands_info: &OperandsInfo, + ) -> Result<()> { + convert_g2s_cta(ctx, rewriter, self.get_operation(), operands_info, 4) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CpAsyncBulkTensorS2gTile4dOp { fn convert( @@ -503,6 +553,18 @@ impl MirToLlvmConversion for CpAsyncBulkTensorG2sTile5dOp { } } +#[op_interface_impl] +impl MirToLlvmConversion for CpAsyncBulkTensorG2sCtaTile5dOp { + fn convert( + &self, + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + operands_info: &OperandsInfo, + ) -> Result<()> { + convert_g2s_cta(ctx, rewriter, self.get_operation(), operands_info, 5) + } +} + #[op_interface_impl] impl MirToLlvmConversion for CpAsyncBulkTensorS2gTile5dOp { fn convert( 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 f385fe8109..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 6d49318480..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 fd4aa60aa9..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 d6822754f5..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// 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 c60e551e5a..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated conversion interfaces: `wgmma_control` intrinsics. @@ -21,7 +21,6 @@ use pliron::context::{Context, Ptr}; use pliron::derive::op_interface_impl; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; 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/asm.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/asm.rs index 908278ed82..36e00d407d 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/asm.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/asm.rs @@ -13,7 +13,6 @@ use pliron::context::{Context, Ptr}; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; @@ -25,7 +24,7 @@ pub(crate) fn convert_inline_ptx( op: Ptr, _operands_info: &OperandsInfo, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let inline_ptx = InlinePtxOp::new(op); let template = inline_ptx .get_attr_ptx_template(ctx) diff --git a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/common.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/common.rs index b766e17e90..6b947c3c40 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/common.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/common.rs @@ -643,7 +643,7 @@ mod tests { #[test] fn inline_asm_convergent_copies_the_source_op_location() { - use pliron::location::{Located, Location}; + use pliron::location::Location; let mut ctx = make_ctx(); let (module_ptr, entry) = build_test_func(&mut ctx, vec![]); @@ -660,7 +660,7 @@ mod tests { let asms = find_body_ops::(&ctx, module_ptr); assert_eq!(asms.len(), 1); - assert_eq!(asms[0].get_operation().deref(&ctx).loc(), expected); + assert_eq!(asms[0].get_operation().deref(&ctx).loc(), &expected); } #[test] diff --git a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/memory.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/memory.rs index f5ae10e694..2d83d61489 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/memory.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/memory.rs @@ -12,7 +12,9 @@ //! boundary, not a Rust pointer-address observation (those genericize first, //! see `convert/ops/cast.rs`). -use crate::convert::intrinsics::common::cast_to_shared_addrspace; +use crate::convert::intrinsics::common::{ + cast_to_cluster_shared_addrspace, cast_to_shared_addrspace, +}; use llvm_export::op_interfaces::CastOpInterface; use llvm_export::ops as llvm; use pliron::builtin::types::{IntegerType, Signedness}; @@ -36,7 +38,59 @@ pub(crate) fn convert_cvta_generic_to_shared_offset( if operands.is_empty() { return pliron::input_err_noloc!("cvta_generic_to_shared_offset requires an operand"); } - let shared_ptr = cast_to_shared_addrspace(ctx, rewriter, operands[0]); + let space = operands[0] + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(|ty| ty.address_space()); + let shared_ptr = match space { + Some(3 | 7) => operands[0], // AS7 retains the remote CTA rank. + Some(0) => cast_to_shared_addrspace(ctx, rewriter, operands[0]), + _ => { + return pliron::input_err_noloc!( + "shared offset requires generic, shared or cluster shared pointer" + ); + } + }; + + let result_ty = op.deref(ctx).get_result(0).get_type(ctx); + let Some(result_width) = result_ty + .deref(ctx) + .downcast_ref::() + .map(IntegerType::width) + else { + return pliron::input_err_noloc!( + "cvta_generic_to_shared_offset requires an integer result" + ); + }; + let result_ty = IntegerType::get(ctx, result_width, Signedness::Signless); + let ptr_to_int = llvm::PtrToIntOp::new(ctx, shared_ptr, result_ty.into()); + rewriter.insert_operation(ctx, ptr_to_int.get_operation()); + rewriter.replace_operation(ctx, op, ptr_to_int.get_operation()); + Ok(()) +} + +pub(crate) fn convert_cvta_generic_to_cluster_shared_offset( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, +) -> Result<()> { + let operands: Vec<_> = op.deref(ctx).operands().collect(); + if operands.is_empty() { + return pliron::input_err_noloc!("cvta_generic_to_shared_offset requires an operand"); + } + let space = operands[0] + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(|ty| ty.address_space()); + if !matches!(space, Some(0 | 3 | 7)) { + return pliron::input_err_noloc!( + "cluster shared offset requires generic or shared pointer" + ); + } + let shared_ptr = cast_to_cluster_shared_addrspace(ctx, rewriter, operands[0]); let result_ty = op.deref(ctx).get_result(0).get_type(ctx); let Some(result_width) = result_ty @@ -54,3 +108,96 @@ pub(crate) fn convert_cvta_generic_to_shared_offset( rewriter.replace_operation(ctx, op, ptr_to_int.get_operation()); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::convert::ops::test_util::{ + append_mir_return, build_kernel, find_all, kernel_blocks, make_ctx, + }; + use dialect_mir::types::MirPtrType; + use dialect_nvvm::ops::{CvtaGenericToClusterSharedOffsetOp, CvtaGenericToSharedOffsetOp}; + use pliron::r#type::TypeHandle; + #[test] + fn cvta_preserves_local_and_cluster_offsets() { + for space in [0, 3, 7] { + for width in [32, 64] { + let mut ctx = make_ctx(); + let word: TypeHandle = IntegerType::get(&ctx, width, Signedness::Unsigned).into(); + let ptr: TypeHandle = MirPtrType::get(&mut ctx, word, true, space).into(); + let (module, block) = build_kernel(&mut ctx, vec![ptr], vec![word]); + let arg = block.deref(&ctx).get_argument(0); + let cvta = Operation::new( + &mut ctx, + CvtaGenericToSharedOffsetOp::get_concrete_op_info(), + vec![word], + vec![arg], + vec![], + 0, + ); + cvta.insert_at_back(block, &ctx); + let value = cvta.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![value]); + crate::lower_mir_to_llvm(&mut ctx, module).unwrap(); + let blocks = kernel_blocks(&ctx, module); + let casts = find_all::(&ctx, &blocks); + let ints = find_all::(&ctx, &blocks); + assert_eq!(casts.len(), usize::from(space == 0)); + assert_eq!(ints.len(), 1); + let operand = ints[0].get_operation().deref(&ctx).get_operand(0); + let ty = operand.get_type(&ctx); + assert_eq!( + ty.deref(&ctx) + .downcast_ref::() + .unwrap() + .address_space(), + if space == 0 { 3 } else { space } + ); + if space == 0 { + assert_eq!(operand, casts[0].get_operation().deref(&ctx).get_result(0)); + } + } + } + } + #[test] + fn cluster_cvta_targets_cluster_window_from_generic_and_named_pointers() { + for space in [0, 3, 7] { + for width in [32, 64] { + let mut ctx = make_ctx(); + let word: TypeHandle = IntegerType::get(&ctx, width, Signedness::Unsigned).into(); + let ptr: TypeHandle = MirPtrType::get(&mut ctx, word, true, space).into(); + let (module, block) = build_kernel(&mut ctx, vec![ptr], vec![word]); + let arg = block.deref(&ctx).get_argument(0); + let cvta = Operation::new( + &mut ctx, + CvtaGenericToClusterSharedOffsetOp::get_concrete_op_info(), + vec![word], + vec![arg], + vec![], + 0, + ); + cvta.insert_at_back(block, &ctx); + let value = cvta.deref(&ctx).get_result(0); + append_mir_return(&mut ctx, block, vec![value]); + crate::lower_mir_to_llvm(&mut ctx, module).unwrap(); + let blocks = kernel_blocks(&ctx, module); + let casts = find_all::(&ctx, &blocks); + let ints = find_all::(&ctx, &blocks); + assert_eq!(casts.len(), usize::from(space != 7)); + assert_eq!(ints.len(), 1); + let operand = ints[0].get_operation().deref(&ctx).get_operand(0); + let ty = operand.get_type(&ctx); + assert_eq!( + ty.deref(&ctx) + .downcast_ref::() + .unwrap() + .address_space(), + 7 + ); + if space != 7 { + assert_eq!(operand, casts[0].get_operation().deref(&ctx).get_result(0)); + } + } + } + } +} diff --git a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/tma.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/tma.rs index 9c5fcae464..66a51c81eb 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/tma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/tma.rs @@ -58,6 +58,166 @@ fn g2s_inline_asm(dims: usize, multicast: bool, cta_group: i32) -> (String, Stri (template, constraints.join(",")) } +fn g2s_cta_inline_asm(dims: usize) -> (String, String) { + let coordinates = (0..dims) + .map(|index| format!("${}", 3 + index)) + .collect::>() + .join(", "); + let template = format!( + "cp.async.bulk.tensor.{dims}d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$0], [$2, {{{coordinates}}}], [$1];" + ); + let mut constraints = vec!["r", "r", "l"]; + constraints.extend(std::iter::repeat_n("r", dims)); + constraints.push("~{memory}"); + (template, constraints.join(",")) +} + +/// Address space of a lowered pointer operand, or `None` if it is not a +/// pointer. +fn pointer_address_space(ctx: &Context, value: pliron::value::Value) -> Option { + value + .get_type(ctx) + .deref(ctx) + .downcast_ref::() + .map(|pointer| pointer.address_space()) +} + +/// Convert a TMA G2S operation whose destination and barrier belong to the +/// issuing CTA rather than to cluster shared memory. +/// +/// The destination is shared by type: the device wrapper takes a +/// `SharedPtr`, which lowers to an `addrspace(3)` pointer, while the raw +/// ABI entry still passes a generic pointer under the caller's promise of +/// CTA locality. The barrier is a caller promise on both layers, as for every +/// mbarrier operation. Generic and shared pointers are converted to the +/// shared window. A pointer known to be cluster shared (`addrspace(7)`, for +/// example from `map_shared_rank`) may name another CTA's shared memory, so +/// it is rejected rather than reinterpreted as a local offset; such copies +/// use the cluster form. The check sees only uncast operands: a pointer cast +/// or `SharedPtr::from_raw` applied to a cluster pointer is the caller's +/// unsafe contract. +pub(crate) fn convert_g2s_cta( + ctx: &mut Context, + rewriter: &mut DialectConversionRewriter, + op: Ptr, + _operands_info: &OperandsInfo, + dims: usize, +) -> Result<()> { + let operands: Vec<_> = op.deref(ctx).operands().collect(); + let expected_operands = 3 + dims + 1; + if operands.len() != expected_operands { + return pliron::input_err_noloc!( + "TMA CTA G2S {}D requires {} operands, got {}", + dims, + expected_operands, + operands.len() + ); + } + for (index, role) in [(0, "destination"), (1, "barrier")] { + match pointer_address_space(ctx, operands[index]) { + Some(0 | 3) => {} + Some(7) => { + return pliron::input_err!( + op.deref(ctx).loc(), + "cp_async_bulk_tensor_{}d_g2s_cta requires a CTA-local shared {}, \ + but it is a cluster shared memory pointer (addrspace 7), which may \ + name another CTA's shared memory; use cp_async_bulk_tensor_{}d_g2s \ + for cluster destinations", + dims, + role, + dims + ); + } + Some(space) => { + return pliron::input_err!( + op.deref(ctx).loc(), + "cp_async_bulk_tensor_{}d_g2s_cta requires a shared memory {}, \ + but it is in address space {}", + dims, + role, + space + ); + } + None => { + return pliron::input_err!( + op.deref(ctx).loc(), + "cp_async_bulk_tensor_{}d_g2s_cta {} must be a pointer", + dims, + role + ); + } + } + } + + let dst_shared = cast_to_shared_addrspace(ctx, rewriter, operands[0]); + let barrier_shared = cast_to_shared_addrspace(ctx, rewriter, operands[1]); + let void_ty = llvm_types::VoidType::get(ctx); + + if context::lowering_options(ctx).intrinsic_backend == IntrinsicBackend::LibNvvm { + // Legacy NVVM IR has no typed CTA-destination intrinsic. Pass the + // 32-bit shared-window addresses to the exact PTX instruction. + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let dst_address = llvm::PtrToIntOp::new(ctx, dst_shared, i32_ty.into()); + rewriter.insert_operation(ctx, dst_address.get_operation()); + let barrier_address = llvm::PtrToIntOp::new(ctx, barrier_shared, i32_ty.into()); + rewriter.insert_operation(ctx, barrier_address.get_operation()); + let mut inputs = vec![ + dst_address.get_operation().deref(ctx).get_result(0), + barrier_address.get_operation().deref(ctx).get_result(0), + operands[2], + ]; + inputs.extend(operands[3..3 + dims].iter().copied()); + let (template, constraints) = g2s_cta_inline_asm(dims); + inline_asm_convergent( + ctx, + rewriter, + op, + void_ty.into(), + inputs, + &template, + &constraints, + ); + rewriter.erase_operation(ctx, op); + return Ok(()); + } + + // LLVM NVPTX: the typed intrinsic the catalog probes record. + let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless); + let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless); + let i1_ty = IntegerType::get(ctx, 1, Signedness::Signless); + let smem_ptr_ty = llvm_types::PointerType::get(ctx, 3); + let generic_ptr_ty = llvm_types::PointerType::get(ctx, 0); + let mut arg_types: Vec = vec![ + smem_ptr_ty.into(), + smem_ptr_ty.into(), + generic_ptr_ty.into(), + ]; + for _ in 0..dims { + arg_types.push(i32_ty.into()); + } + arg_types.push(i64_ty.into()); // cache_hint + arg_types.push(i1_ty.into()); // use_cache_hint + + let intrinsic_name = format!("llvm_nvvm_cp_async_bulk_tensor_g2s_cta_tile_{}d", dims); + let func_ty = llvm_types::FuncType::get(ctx, void_ty.into(), arg_types, false); + let parent_block = op.deref(ctx).get_parent_block().unwrap(); + helpers::ensure_intrinsic_declared(ctx, parent_block, &intrinsic_name, func_ty) + .map_err(|e| pliron::input_error_noloc!("{}", e))?; + + let mut call_args = vec![dst_shared, barrier_shared]; + call_args.extend(operands[2..].iter().copied()); + let use_cache_hint = create_i1_const(ctx, rewriter, false); + call_args.push(use_cache_hint); + + let sym_name: pliron::identifier::Identifier = intrinsic_name.as_str().try_into().unwrap(); + let callee = CallOpCallable::Direct(sym_name); + let llvm_call = llvm::CallOp::new(ctx, callee, func_ty, call_args); + crate::convert::preserve_location(ctx, op, llvm_call.get_operation()); + rewriter.insert_operation(ctx, llvm_call.get_operation()); + rewriter.erase_operation(ctx, op); + Ok(()) +} + fn convert_g2s_impl( ctx: &mut Context, rewriter: &mut DialectConversionRewriter, @@ -746,7 +906,7 @@ pub(crate) fn convert_control( #[cfg(test)] mod tests { - use super::{g2s_inline_asm, reduce_inline_asm, s2g_inline_asm}; + use super::{g2s_cta_inline_asm, g2s_inline_asm, reduce_inline_asm, s2g_inline_asm}; #[test] fn inline_tma_templates_keep_exact_ptx_shapes() { @@ -764,6 +924,13 @@ mod tests { "l,l,l,r,r,h,~{memory}".into(), ) ); + assert_eq!( + g2s_cta_inline_asm(2), + ( + "cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$0], [$2, {$3, $4}], [$1];".into(), + "r,r,l,r,r,~{memory}".into(), + ) + ); assert_eq!( s2g_inline_asm(5), ( diff --git a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/wgmma.rs b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/wgmma.rs index 29eeb935f8..0f76c7a81a 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/intrinsics/wgmma.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/intrinsics/wgmma.rs @@ -16,7 +16,6 @@ use pliron::context::{Context, Ptr}; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; @@ -451,7 +450,7 @@ fn convert_mma_group_values_for_shape( where F: FnOnce(usize) -> String, { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let result_count = op.deref(ctx).get_num_results(); let operands: Vec<_> = op.deref(ctx).operands().collect(); @@ -541,7 +540,7 @@ fn convert_mma_loop_values_for_kind( op: Ptr, input_kind: WgmmaInputKind, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let result_count = op.deref(ctx).get_num_results(); let operands: Vec<_> = op.deref(ctx).operands().collect(); @@ -604,7 +603,7 @@ pub(crate) fn convert_mma_loop_pipeline_values( op: Ptr, _operands_info: &OperandsInfo, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let result_count = op.deref(ctx).get_num_results(); let operands: Vec<_> = op.deref(ctx).operands().collect(); @@ -680,7 +679,7 @@ pub(crate) fn convert_mma_pipeline_values( op: Ptr, _operands_info: &OperandsInfo, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let result_count = op.deref(ctx).get_num_results(); let operands: Vec<_> = op.deref(ctx).operands().collect(); diff --git a/cuda-oxide/crates/mir-lower/src/convert/mod.rs b/cuda-oxide/crates/mir-lower/src/convert/mod.rs index d4f028e537..fd0f44c956 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/mod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/mod.rs @@ -28,7 +28,6 @@ use pliron::{ context::{Context, Ptr}, - location::Located, operation::Operation, }; @@ -38,7 +37,9 @@ pub(crate) fn preserve_location( source: Ptr, lowered: Ptr, ) -> Ptr { - lowered.deref_mut(ctx).set_loc(source.deref(ctx).loc()); + lowered + .deref_mut(ctx) + .set_loc(source.deref(ctx).loc().clone()); lowered } diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/addressing.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/addressing.rs index 8e5670592b..2f4da18622 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/addressing.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/addressing.rs @@ -7,7 +7,7 @@ use super::common::{anyhow_to_pliron, byte_offset_gep}; use crate::convert::enum_payload_storage::enum_payload_storage_type; use crate::convert::types::{ StructLayoutInfo, build_enum_slot_map, build_struct_slot_map, convert_type, - llvm_type_size_align, mir_element_stride, mir_type_abi_align, + llvm_type_size_align, mir_element_stride, mir_type_abi_align, union_field_storage_type, }; use dialect_mir::ops::{MirConstantOp, MirFieldAddrOp}; use dialect_mir::types::{MirEnumType, MirPtrType, MirStructType, MirTupleType, MirUnionType}; @@ -18,7 +18,6 @@ use pliron::context::{Context, Ptr}; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; @@ -74,18 +73,52 @@ pub(crate) fn convert_field_addr( } }; - let union_field_count = mir_ptr_pointee + let union_field = mir_ptr_pointee .deref(ctx) .downcast_ref::() - .map(MirUnionType::field_count); - if let Some(field_count) = union_field_count { - if field_index >= field_count { + .map(|union_ty| { + ( + union_ty.name().to_string(), + union_ty.field_count(), + union_ty.get_field_type(field_index), + ) + }); + if let Some((union_name, field_count, field_ty)) = union_field { + let Some(field_ty) = field_ty else { return pliron::input_err_noloc!( "field_addr index {} out of bounds for union with {} fields", field_index, field_count ); + }; + + // A union field's bytes hold its STORAGE type, the type the bytes + // really have, which can differ from its SEMANTIC type, the type the + // program sees. A field holding a shared-memory pointer is stored as + // a generic pointer so the union's byte image is the same on every + // target (a shared pointer is 4 bytes under modern NVVM's p3:32 and + // 8 bytes elsewhere); the value paths (insert/extract) convert at + // that boundary. The address formed here ESCAPES to loads and stores + // typed with the semantic type, which the storage bytes cannot honor: + // a raw shared pointer stored through it is a 4-byte write into the + // 8-byte slot under p3:32, and a load through it reads the generic + // address back as a shared offset on the other paths. Fail closed, + // as the enum arm below does for payload fields. Fields whose storage + // is their semantic type (integers, byte arrays, generic pointers) + // keep the address path. The diagnostic calls the storage type + // "canonical", the wording the enum arm and its fixture share. + let semantic_ty = convert_type(ctx, field_ty).map_err(anyhow_to_pliron)?; + let storage_ty = union_field_storage_type(ctx, semantic_ty).map_err(anyhow_to_pliron)?; + if storage_ty != semantic_ty { + return pliron::input_err_noloc!( + "field_addr: cannot hand out the in-place address of union field {} of `{}`: its bytes use canonical storage type {} while its semantic type is {}, and loads or stores made through an escaped field address are typed with the semantic type, which the canonical bytes cannot honor; build the union by value and read the field by value; in-place writes and borrows of this field are refused", + field_index, + union_name, + storage_ty.deref(ctx).disp(ctx), + semantic_ty.deref(ctx).disp(ctx) + ); } + // Every union field begins at byte zero. Emit an explicit zero-offset // GEP instead of forwarding the base SSA value directly: the distinct // result keeps dialect conversion's pointer-type history unambiguous @@ -373,7 +406,7 @@ pub(crate) fn convert_array_element_addr( op: Ptr, _operands_info: &OperandsInfo, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let arr_ptr = op.deref(ctx).get_operand(0); let index = op.deref(ctx).get_operand(1); @@ -1039,6 +1072,100 @@ mod tests { ); } + /// The union arm of the same gate. A union field that holds a + /// shared-memory pointer is stored as a GENERIC pointer, so an in-place + /// address would let a semantic `ptr addrspace(3)` store write 4 bytes + /// under modern NVVM's `p3:32` into the 8-byte slot, and would let a + /// load read the generic address back as a shared offset elsewhere. + #[test] + fn shared_pointer_union_field_addr_fails_closed() { + let mut ctx = make_ctx(); + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); + let bits: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let union_ty: TypeHandle = MirUnionType::get( + &mut ctx, + "SharedPtrBits".into(), + vec!["ptr".into(), "bits".into()], + vec![shared, bits], + 8, + 8, + ) + .into(); + let union_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, union_ty, true).into(); + let field_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, shared, true).into(); + + let (module, block) = build_kernel(&mut ctx, vec![union_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + let op = MirFieldAddrOp::build(&mut ctx, base, field_ptr_ty, 0).expect("field_addr build"); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + let err = crate::lower_mir_to_llvm(&mut ctx, module) + .expect_err("addressing a shared-pointer union field in place must fail to lower"); + assert!( + err.err.to_string().contains("canonical storage type"), + "unexpected error: {}", + err.err + ); + } + + /// A union field whose storage IS its semantic type keeps the address + /// path: one zero-offset `i8` GEP off the ORIGINAL union pointer, no + /// stack spill. This is the `union_aggregate` write form + /// (`bits.bytes = [..]` on a `{ word: u32, bytes: [u8; 4] }` union). + #[test] + fn plain_union_field_addr_lowers_to_zero_offset_gep() { + use llvm_export::ops::GepIndex; + + let mut ctx = make_ctx(); + let u8_ty: TypeHandle = IntegerType::get(&ctx, 8, Signedness::Unsigned).into(); + let u32_ty: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let bytes_ty: TypeHandle = MirArrayType::get(&mut ctx, u8_ty, 4).into(); + let union_ty: TypeHandle = MirUnionType::get( + &mut ctx, + "Bits".into(), + vec!["word".into(), "bytes".into()], + vec![u32_ty, bytes_ty], + 4, + 4, + ) + .into(); + let union_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, union_ty, true).into(); + let bytes_ptr_ty: TypeHandle = MirPtrType::get_generic(&mut ctx, bytes_ty, true).into(); + + let (module, block) = build_kernel(&mut ctx, vec![union_ptr_ty], vec![]); + let base = block.deref(&ctx).get_argument(0); + let op = MirFieldAddrOp::build(&mut ctx, base, bytes_ptr_ty, 1).expect("field_addr build"); + op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "a plain union field address must not spill the union to a stack copy" + ); + let geps = find_all::(&ctx, &body); + assert_eq!(geps.len(), 1, "one union field_addr lowers to one GEP"); + let gep = &geps[0]; + assert_eq!( + gep.get_operation().deref(&ctx).get_operand(0), + base, + "the GEP must address the ORIGINAL union storage" + ); + assert_eq!( + gep.src_elem_type(&ctx), + IntegerType::get(&ctx, 8, Signedness::Signless).into(), + "every union field starts at byte zero, addressed in byte units" + ); + assert!( + matches!(gep.indices(&ctx).as_slice(), [GepIndex::Constant(0)]), + "a union field address is a zero-offset GEP" + ); + } + /// A payload with no slot of its own is addressed at its byte offset off /// the ORIGINAL enum pointer: no stack spill is introduced, so a write /// through the result lands in the enum rather than in a copy. diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs index 41ca9756cf..e60a2878d1 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs @@ -4,9 +4,10 @@ */ use super::common::{anyhow_to_pliron, spill_enum_value}; +use crate::convert::target_stable_storage::coerce_target_stable_value; use crate::convert::types::{ StructLayoutInfo, StructSlotMap, build_struct_slot_map, build_union_storage_type, convert_type, - is_zero_sized_type, make_slice_struct, + is_zero_sized_type, make_slice_struct, union_field_storage_type, }; use dialect_mir::ops::{MirExtractFieldOp, MirInsertFieldOp}; use dialect_mir::types::{ @@ -311,10 +312,15 @@ fn convert_extract_union_field( let storage_ty = build_union_storage_type(ctx, &union_ty).map_err(anyhow_to_pliron)?; let ptr = spill_enum_value(ctx, rewriter, union_value, storage_ty, union_ty.abi_align()); - let load = llvm::LoadOp::new(ctx, ptr, field_llvm_ty); + let field_storage_ty = + union_field_storage_type(ctx, field_llvm_ty).map_err(anyhow_to_pliron)?; + let load = llvm::LoadOp::new(ctx, ptr, field_storage_ty); llvm_export::ops::set_op_alignment(ctx, load.get_operation(), union_ty.abi_align() as u32); rewriter.insert_operation(ctx, load.get_operation()); - rewriter.replace_operation(ctx, op, load.get_operation()); + let stored = load.get_operation().deref(ctx).get_result(0); + let semantic = + coerce_target_stable_value(ctx, rewriter, stored, field_llvm_ty, "union field storage")?; + rewriter.replace_operation_with_values(ctx, op, vec![semantic]); Ok(()) } @@ -344,7 +350,16 @@ fn convert_insert_union_field( let storage_ty = build_union_storage_type(ctx, &union_ty).map_err(anyhow_to_pliron)?; let ptr = spill_enum_value(ctx, rewriter, union_value, storage_ty, union_ty.abi_align()); - let store = llvm::StoreOp::new(ctx, new_value, ptr); + let field_storage_ty = + union_field_storage_type(ctx, field_llvm_ty).map_err(anyhow_to_pliron)?; + let stored_value = coerce_target_stable_value( + ctx, + rewriter, + new_value, + field_storage_ty, + "union field storage", + )?; + let store = llvm::StoreOp::new(ctx, stored_value, ptr); llvm_export::ops::set_op_alignment(ctx, store.get_operation(), union_ty.abi_align() as u32); rewriter.insert_operation(ctx, store.get_operation()); @@ -363,6 +378,7 @@ mod tests { use crate::convert::ops::test_util::*; use dialect_mir::attributes::FieldIndexAttr; use dialect_mir::ops as mir; + use dialect_mir::types::MirPtrType; use llvm_export::types as llvm_types; @@ -431,6 +447,176 @@ mod tests { ); } + /// A shared pointer carried through a union must use generic-pointer + /// physical storage in both directions. Construction converts AS3 -> AS0; + /// extraction converts the stored AS0 pointer back to semantic AS3. + #[test] + fn shared_pointer_union_field_round_trips_through_generic_storage() { + let mut ctx = make_ctx(); + + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); + let bits: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let union_ty: TypeHandle = MirUnionType::get( + &mut ctx, + "SharedPointerBits".into(), + vec!["ptr".into(), "bits".into()], + vec![shared, bits], + 8, + 8, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![shared], vec![shared]); + let pointer = block.deref(&ctx).get_argument(0); + + let undef = mir::MirUndefOp::new(&mut ctx, union_ty).get_operation(); + undef.insert_at_back(block, &ctx); + let undef_value = undef.deref(&ctx).get_result(0); + + let insert = Operation::new( + &mut ctx, + MirInsertFieldOp::get_concrete_op_info(), + vec![union_ty], + vec![undef_value, pointer], + vec![], + 0, + ); + MirInsertFieldOp::new(insert).set_attr_insert_index(&ctx, FieldIndexAttr(0)); + insert.insert_at_back(block, &ctx); + let union_value = insert.deref(&ctx).get_result(0); + + let extract = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![shared], + vec![union_value], + vec![], + 0, + ); + MirExtractFieldOp::new(extract).set_attr_index(&ctx, FieldIndexAttr(0)); + extract.insert_at_back(block, &ctx); + let result = extract.deref(&ctx).get_result(0); + + append_mir_return(&mut ctx, block, vec![result]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + + assert_eq!( + count_ops::(&ctx, &body), + 2, + "union insertion must genericize AS3 and extraction must restore AS3" + ); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "MirInsertFieldOp must be fully lowered" + ); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "MirExtractFieldOp must be fully lowered" + ); + } + + #[test] + fn nested_shared_pointer_union_field_round_trips_through_generic_storage() { + let mut ctx = make_ctx(); + + let pointee: TypeHandle = IntegerType::get(&ctx, 32, Signedness::Unsigned).into(); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, pointee, true).into(); + let cookie: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let inner: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "SharedPointerInner".into(), + vec!["pointer".into(), "cookie".into()], + vec![shared, cookie], + vec![0, 1], + vec![0, 8], + 16, + 8, + ) + .into(); + let union_ty: TypeHandle = MirUnionType::get( + &mut ctx, + "NestedSharedPointerUnion".into(), + vec!["inner".into()], + vec![inner], + 16, + 8, + ) + .into(); + + let (module, block) = build_kernel(&mut ctx, vec![shared, cookie], vec![shared]); + let pointer = block.deref(&ctx).get_argument(0); + let cookie_value = block.deref(&ctx).get_argument(1); + + let construct_inner = Operation::new( + &mut ctx, + mir::MirConstructStructOp::get_concrete_op_info(), + vec![inner], + vec![pointer, cookie_value], + vec![], + 0, + ); + construct_inner.insert_at_back(block, &ctx); + let inner_value = construct_inner.deref(&ctx).get_result(0); + + let undef = mir::MirUndefOp::new(&mut ctx, union_ty).get_operation(); + undef.insert_at_back(block, &ctx); + let undef_value = undef.deref(&ctx).get_result(0); + + let insert = Operation::new( + &mut ctx, + MirInsertFieldOp::get_concrete_op_info(), + vec![union_ty], + vec![undef_value, inner_value], + vec![], + 0, + ); + MirInsertFieldOp::new(insert).set_attr_insert_index(&ctx, FieldIndexAttr(0)); + insert.insert_at_back(block, &ctx); + let union_value = insert.deref(&ctx).get_result(0); + + let extract_inner = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![inner], + vec![union_value], + vec![], + 0, + ); + MirExtractFieldOp::new(extract_inner).set_attr_index(&ctx, FieldIndexAttr(0)); + extract_inner.insert_at_back(block, &ctx); + let recovered_inner = extract_inner.deref(&ctx).get_result(0); + + let extract_pointer = Operation::new( + &mut ctx, + MirExtractFieldOp::get_concrete_op_info(), + vec![shared], + vec![recovered_inner], + vec![], + 0, + ); + MirExtractFieldOp::new(extract_pointer).set_attr_index(&ctx, FieldIndexAttr(0)); + extract_pointer.insert_at_back(block, &ctx); + let result = extract_pointer.deref(&ctx).get_result(0); + + append_mir_return(&mut ctx, block, vec![result]); + + crate::lower_mir_to_llvm(&mut ctx, module).expect("lowering failed"); + let body = kernel_blocks(&ctx, module); + + assert_eq!( + count_ops::(&ctx, &body), + 2, + "nested union insertion must genericize AS3 and extraction must restore AS3" + ); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!(count_ops::(&ctx, &body), 0); + } + /// The row-width arm of [`resolve_aggregate_slots`]' no-history fallback /// is defensive symmetry: no current lowering path produces a runtime-width /// slice value with an empty conversion history, so no end-to-end pipeline diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/test_support.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/test_support.rs index 22d3a22d4a..bc66f12485 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/test_support.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/test_support.rs @@ -18,6 +18,8 @@ use pliron::operation::Operation; use pliron::r#type::TypeHandle; use pliron::value::Value; +use crate::convert::types::MAX_NATURAL_ALIGN; + pub(super) fn insert_indices(ctx: &Context, inserts: &[llvm::InsertValueOp]) -> Vec> { inserts.iter().map(|op| op.indices(ctx)).collect() } @@ -82,7 +84,7 @@ pub(super) fn unit_niche_enum( .collect::>(); let discriminants = (0..inhabited.len() as u64).collect(); let carrier_size = u64::from(carrier_width).div_ceil(8); - let carrier_align = carrier_size.next_power_of_two().min(16); + let carrier_align = carrier_size.next_power_of_two().min(MAX_NATURAL_ALIGN); MirEnumType::get_with_encoding( ctx, "UnitNiche".into(), diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/arithmetic.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/arithmetic.rs index 0c16f6f7e0..0ce5d1d8f9 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/arithmetic.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/arithmetic.rs @@ -44,7 +44,6 @@ use pliron::context::{Context, Ptr}; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; @@ -57,7 +56,7 @@ use pliron::value::Value; /// Extract binary operands from the (already-converted) operation. fn get_binary_operands(op: Ptr, ctx: &Context) -> Result<(Value, Value)> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let operands: Vec<_> = op.deref(ctx).operands().collect(); match operands.as_slice() { [lhs, rhs] => Ok((*lhs, *rhs)), @@ -350,7 +349,7 @@ fn convert_checked_binop_with_intrinsic( rhs: Value, op_name: &str, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let lhs_ty = lhs.get_type(ctx); let width = lhs_ty @@ -751,7 +750,7 @@ pub(crate) fn convert_three_way_cmp( // and f32::total_cmp lowers to integer bit-twiddling, not Cmp. // An OLT/OGT select chain would also miscompile: NaN compares // false on both predicates and would silently yield Equal. - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); return pliron::input_err!( loc, "BinOp::Cmp on floats is never emitted by rustc; total-order lowering unimplemented" @@ -763,7 +762,7 @@ pub(crate) fn convert_three_way_cmp( let is_gt = emit_cmp_value(ctx, rewriter, op, lhs, rhs, is_signed, false); let mir_result_ty = op.deref(ctx).get_result(0).get_type(ctx); - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let (discr_ty, variant_discriminants) = { let mir_result_ty_obj = mir_result_ty.deref(ctx); let enum_ty = mir_result_ty_obj @@ -834,7 +833,7 @@ fn emit_cmp_value( (false, false) => ICmpPredicateAttr::UGT, }; let cmp_op = llvm::ICmpOp::new(ctx, pred, lhs, rhs).get_operation(); - cmp_op.deref_mut(ctx).set_loc(op.deref(ctx).loc()); + cmp_op.deref_mut(ctx).set_loc(op.deref(ctx).loc().clone()); rewriter.insert_operation(ctx, cmp_op); cmp_op.deref(ctx).get_result(0) } diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs index 3eccb40fcf..bc19168b0c 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs @@ -90,7 +90,6 @@ use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInf use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; use pliron::linked_list::ContainsLinkedList; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; @@ -738,7 +737,7 @@ pub fn convert( if resolved_name.starts_with("__nv_") && let Some(parent_block) = parent_block { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); helpers::ensure_intrinsic_declared(ctx, parent_block, &resolved_name, func_type) .map_err(|e| { pliron::input_error!( @@ -808,7 +807,7 @@ pub fn convert( // dropped a real value; report that instead of letting pliron's // erase-with-uses invariant panic and escape into rustc as an ICE. if op.deref(ctx).has_use() { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); return pliron::input_err!( loc, "mir.call lowering produced a void LLVM call but the MIR op \ @@ -837,7 +836,7 @@ fn convert_rust_select_unpredictable( rewriter: &mut DialectConversionRewriter, op: Ptr, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); if op.deref(ctx).get_num_results() != 1 { return pliron::input_err!(loc, "select_unpredictable call must have one result"); } @@ -866,7 +865,7 @@ fn convert_rust_bit_intrinsic( op: Ptr, intrinsic: RustBitIntrinsic, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); if op.deref(ctx).get_num_results() != 1 { return pliron::input_err!(loc, "Rust bit intrinsic call must have one result"); } @@ -993,7 +992,7 @@ fn convert_rust_saturating_intrinsic( operands_info: &OperandsInfo, intrinsic: RustSaturatingIntrinsic, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); if op.deref(ctx).get_num_results() != 1 { return pliron::input_err!(loc, "Rust saturating intrinsic call must have one result"); } @@ -1074,7 +1073,7 @@ fn convert_rust_exact_div( op: Ptr, operands_info: &OperandsInfo, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); if op.deref(ctx).get_num_results() != 1 { return pliron::input_err!(loc, "exact_div call must have one result"); } @@ -1150,7 +1149,7 @@ fn convert_rust_carrying_mul_add( op: Ptr, operands_info: &OperandsInfo, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); if op.deref(ctx).get_num_results() != 1 { return pliron::input_err!(loc, "Rust carrying_mul_add intrinsic must have one result"); } @@ -1272,7 +1271,7 @@ fn convert_rust_float_math_intrinsic( op: Ptr, intrinsic: RustFloatMathIntrinsic, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); if op.deref(ctx).get_num_results() != 1 { return pliron::input_err!(loc, "Rust float math intrinsic call must have one result"); } 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 26de49c018..6d9aff6b90 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs @@ -70,7 +70,6 @@ use pliron::context::{Context, Ptr}; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::printable::Printable; @@ -87,7 +86,7 @@ pub fn convert( op: Ptr, operands_info: &OperandsInfo, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let operands: Vec<_> = op.deref(ctx).operands().collect(); let val = match operands.as_slice() { @@ -1744,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); @@ -1923,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-lower/src/convert/ops/constants.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/constants.rs index f2ec4bbaca..e6292f2b80 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/constants.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/constants.rs @@ -30,7 +30,6 @@ use pliron::context::{Context, Ptr}; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs index 4e4f899cf4..bf31f820cf 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs @@ -331,7 +331,7 @@ mod tests { use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::Context; use pliron::linked_list::ContainsLinkedList; - use pliron::location::{Located, Location, Source}; + use pliron::location::{Location, Source}; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; @@ -391,7 +391,7 @@ mod tests { let body = kernel_blocks(&ctx, module_ptr); let ret = find_first::(&ctx, &body).expect("expected llvm.return"); - assert_eq!(ret.get_operation().deref(&ctx).loc(), expected); + assert_eq!(ret.get_operation().deref(&ctx).loc(), &expected); } #[test] @@ -626,7 +626,7 @@ mod tests { let (module_ptr, entry) = build_kernel(&mut ctx, vec![i1_ty], vec![]); let cond = entry.deref(&ctx).get_argument(0); let assert_loc = Location::SrcPos { - src: Source::new_from_file(&mut ctx, PathBuf::from("kernel.rs")), + src: Source::new_from_file(&ctx, PathBuf::from("kernel.rs")), pos: combine::stream::position::SourcePosition { line: 44, column: 36, @@ -661,7 +661,7 @@ mod tests { let llvm_br = find_first::(&ctx, &body).expect("expected llvm.cond_br"); assert_eq!( llvm_br.get_operation().deref(&ctx).loc(), - assert_loc, + &assert_loc, "assert branch must keep the MIR assertion source location" ); let false_succ = llvm_br.get_operation().deref(&ctx).get_successor(1); @@ -688,7 +688,7 @@ mod tests { ); assert_eq!( abort_ops[0].deref(&ctx).loc(), - assert_loc, + &assert_loc, "trap call must keep the MIR assertion source location" ); assert!( @@ -697,7 +697,7 @@ mod tests { ); assert_eq!( abort_ops[1].deref(&ctx).loc(), - assert_loc, + &assert_loc, "unreachable must keep the MIR assertion source location" ); @@ -709,7 +709,7 @@ mod tests { .expect("llvm.trap must be declared in the module"); assert_eq!( trap_decl.get_operation().deref(&ctx).loc(), - Location::Unknown, + &Location::Unknown, "the shared llvm.trap declaration must remain locationless" ); } diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/access.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/access.rs index d5b6002ad3..a9b86962ac 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/access.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/access.rs @@ -22,7 +22,6 @@ use pliron::context::{Context, Ptr}; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; @@ -276,7 +275,7 @@ pub(crate) fn convert_ptr_offset( op: Ptr, _operands_info: &OperandsInfo, ) -> Result<()> { - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); let operands: Vec<_> = op.deref(ctx).operands().collect(); let (ptr, offset) = match operands.as_slice() { diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/debug.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/debug.rs index 353ea39e8c..3efbf06519 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/debug.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/debug.rs @@ -10,7 +10,6 @@ use pliron::context::{Context, Ptr}; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/device_global.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/device_global.rs index 4374953b83..f9bd074448 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/device_global.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/device_global.rs @@ -257,6 +257,14 @@ fn create_device_global( }; global_op.set_address_space(ctx, spec.addr_space); global_op.set_source_global_key(ctx, spec.key); + if spec.addr_space == llvm_export::types::address_space::CONSTANT { + // `#[constant]` storage is written by the host after module load. Keep + // the symbol rooted (`@llvm.used`) even when every read of it is + // optimized away, so the host's `cuModuleGetGlobal` still resolves it, + // and stop optimizers from treating its initializer as the value read. + global_op.mark_retained(ctx); + global_op.mark_host_written(ctx); + } if matches!( spec.addr_space, llvm_export::types::address_space::GLOBAL | llvm_export::types::address_space::CONSTANT diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/shared.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/shared.rs index fc413d25b5..1a2503f252 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/shared.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/shared.rs @@ -9,7 +9,7 @@ use super::common::{anyhow_to_pliron, counter_named_global}; use crate::context::{ SharedGlobalDeclaration, SharedGlobalKind, SharedGlobalRecord, SharedGlobalsMap, }; -use crate::convert::types::{convert_type, llvm_type_size_align}; +use crate::convert::types::{convert_type, llvm_type_size_align, required_storage_align}; use crate::helpers; use llvm_export::ops as llvm; use llvm_export::ops::GlobalOpExt; @@ -96,6 +96,22 @@ pub fn convert_shared_alloc_dc( // optional debug identity against the lowered element layout before it // can reach the declaration record or the emitted global. let llvm_elem_type = convert_type(ctx, mir_elem_type).map_err(anyhow_to_pliron)?; + + // Storage is never declared below the element's ABI alignment: a smaller + // request is raised, a larger one is honored. The verifier has already + // rejected a request that is not a power of two. + let required_alignment = required_storage_align(ctx, mir_elem_type, llvm_elem_type) + .ok_or_else(|| { + let elem_display = mir_elem_type.deref(ctx).disp(ctx).to_string(); + pliron::create_error!( + op.deref(ctx).loc(), + pliron::result::ErrorKind::VerificationFailed, + "mir.shared_alloc element type `{elem_display}` has no known ABI alignment; \ + refusing to guess the storage alignment" + ) + })?; + let alignment = alignment.max(required_alignment); + let debug_info = debug_info .filter(|info| shared_debug_type_matches_physical(ctx, info, llvm_elem_type, size)); let debug_owner_function = debug_info.as_ref().and(debug_owner_function); diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/debug_provenance.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/debug_provenance.rs index 5b34793741..02982ba0b5 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/debug_provenance.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/debug_provenance.rs @@ -86,7 +86,7 @@ fn convert_dbg_value_lowers_to_llvm_dbg_value() { .expect("expected lowered llvm.dbg_value marker"); assert_eq!( dbg_value.get_operation().deref(&ctx).loc(), - dbg_loc, + &dbg_loc, "dbg.value lowering should keep the current-value source location" ); let info = llvm::debug_local_variable(&ctx, dbg_value.get_operation()) @@ -285,7 +285,7 @@ fn mem2reg_salvages_tagged_alloca_into_mir_dbg_value() { ); assert_eq!( dbg_values[0].get_operation().deref(&ctx).loc(), - decl_loc, + &decl_loc, "debug records for source-less promoted ops should fall back to the local declaration" ); } diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs index 562db062e5..b245546638 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs @@ -200,6 +200,12 @@ fn convert_global_alloc_places_in_global_or_constant_addrspace() { .starts_with("__device_global_"), "ordinary device globals get the __device_global_ prefix" ); + // The host writes constant memory after load: it stays rooted and must not + // be folded to its initializer. Private device globals stay ordinary. + assert!(global_addr_const.is_retained(&ctx)); + assert!(global_addr_const.is_host_written(&ctx)); + assert!(!global_addr_global.is_retained(&ctx)); + assert!(!global_addr_global.is_host_written(&ctx)); } #[test] diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/enums.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/enums.rs index 5658a39cc1..884c68169a 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/enums.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/enums.rs @@ -44,6 +44,148 @@ fn unit_variants(n: usize) -> Vec { (0..n).map(|i| EnumVariant::unit(format!("V{i}"))).collect() } +/// The alignment anchor is `[0 x iN]` with `N = 8 * align` bits, so an +/// unbounded anchor failed in two ways: from `repr(align(2^21))` up the +/// width passed LLVM's 2^23-bit integer limit (`i8388608` is the last width +/// LLVM accepts, `i8388609` the first it rejects), and at `repr(align(2^29))` +/// the `u32` conversion wrapped to `i0`. `alignment_anchor_type` caps the +/// anchor at `MAX_NATURAL_ALIGN` (16 bytes), which is where cuda-oxide's data +/// layouts and layout model stop anyway; memory operations carry the stronger +/// Rust ABI alignment explicitly. The 8 and 16 rows are the controls, 2^20 is +/// the last width LLVM accepted, 2^21 the first it rejected, 2^29 the wrap. +#[test] +fn enum_conversion_handles_repr_align_bitwidth_boundary() { + use crate::convert::types::{llvm_type_size_align, mir_type_abi_align}; + + for (align, anchor_bits) in [ + (8u64, 64), + (16, 128), + (32, 128), + (64, 128), + (1 << 20, 128), + (1 << 21, 128), + (1 << 29, 128), + ] { + let mut ctx = make_ctx(); + let ty = make_enum_ty(&mut ctx, "AlignedEnum", 8, unit_variants(2), align, align); + let conv = convert_type(&mut ctx, ty).unwrap(); + let array_ty: TypeHandle = MirArrayType::get(&mut ctx, ty, 3).into(); + let conv_array = convert_type(&mut ctx, array_ty).unwrap(); + let natural_align = u64::from(anchor_bits / 8); + + assert_eq!( + llvm_type_size_align(&ctx, conv), + Some((align, natural_align)), + "enum with {align}-byte alignment keeps rustc size" + ); + assert_eq!( + llvm_type_size_align(&ctx, conv_array), + Some((align * 3, natural_align)), + "enum array with {align}-byte alignment keeps rustc stride" + ); + assert_eq!(mir_type_abi_align(&ctx, ty), Some(align)); + assert_eq!(mir_type_abi_align(&ctx, array_ty), Some(align)); + + let conv_ref = conv.deref(&ctx); + let struct_ty = conv_ref + .downcast_ref::() + .expect("converted enum is a struct"); + let anchor_ty = struct_ty + .fields() + .next() + .expect("expected alignment anchor"); + let anchor_ref = anchor_ty.deref(&ctx); + let anchor = anchor_ref + .downcast_ref::() + .expect("alignment anchor is an array"); + assert_eq!(anchor.size(), 0, "alignment anchor must occupy no bytes"); + let anchor_int = anchor.elem_type(); + let anchor_int_ref = anchor_int.deref(&ctx); + assert_eq!( + anchor_int_ref + .downcast_ref::() + .expect("alignment anchor element is an integer") + .width(), + anchor_bits, + "enum with {align}-byte alignment uses a bounded integer anchor" + ); + } +} + +/// Lowering a reference to an over-aligned enum array keeps the Rust ABI +/// alignment on the alloca and the store, while the storage type keeps +/// rustc's size and stride at the capped 16-byte natural alignment. Enum +/// sibling of `convert_ref_preserves_over_aligned_union_array_layout_and_alignment` +/// in `access.rs`; 2^29 is the alignment whose anchor used to wrap to `i0`. +#[test] +fn convert_ref_preserves_over_aligned_enum_array_layout_and_alignment() { + for abi_align in [32u64, 64, 1 << 29] { + let mut ctx = make_ctx(); + let enum_ty = make_enum_ty( + &mut ctx, + "AlignedEnum", + 8, + unit_variants(2), + abi_align, + abi_align, + ); + let array_ty: TypeHandle = MirArrayType::get(&mut ctx, enum_ty, 3).into(); + let mir_ptr_ty = + MirPtrType::get_generic_with_kind(&mut ctx, array_ty, false, MirPointerKind::SharedRef); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + + let undef = mir::MirUndefOp::new(&mut ctx, array_ty); + undef.get_operation().insert_at_back(block, &ctx); + let value = undef.get_operation().deref(&ctx).get_result(0); + let ref_op = Operation::new( + &mut ctx, + mir::MirRefOp::get_concrete_op_info(), + vec![mir_ptr_ty.into()], + vec![value], + vec![], + 0, + ); + let mir_ref = mir::MirRefOp::new(ref_op); + mir_ref.set_mutable(&mut ctx, false); + mir_ref.set_pointer_kind_authority(&mut ctx, MirPointerKindAuthorityAttr::Reborrow); + ref_op.insert_at_back(block, &ctx); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let body = kernel_blocks(&ctx, module_ptr); + let alloca = find_first::(&ctx, &body).expect("expected llvm.alloca"); + let store = find_first::(&ctx, &body).expect("expected llvm.store"); + let llvm_array_ty = alloca.result_pointee_type(&ctx); + let llvm_array_data = llvm_array_ty.deref(&ctx); + let llvm_array = llvm_array_data + .downcast_ref::() + .expect("over-aligned enum array must remain an LLVM array"); + + assert_eq!(llvm_array.size(), 3); + assert_eq!( + crate::convert::types::llvm_type_size_align(&ctx, llvm_array.elem_type()), + Some((abi_align, 16)), + "enum aligned {abi_align} keeps rustc's size at the capped natural alignment" + ); + assert_eq!( + crate::convert::types::llvm_type_size_align(&ctx, llvm_array_ty), + Some((abi_align * 3, 16)), + "enum array aligned {abi_align} keeps rustc's stride" + ); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, alloca.get_operation()), + Some(abi_align as u32), + "alloca of an enum array aligned {abi_align} carries the Rust alignment" + ); + assert_eq!( + llvm_export::ops::op_alignment(&ctx, store.get_operation()), + Some(abi_align as u32), + "store of an enum array aligned {abi_align} carries the Rust alignment" + ); + } +} + /// Converted enum allocation size must equal rustc's `total_size` for /// every memory-faithful tag shape: that size is what GEP strides by. #[test] diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/mod.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/mod.rs index d8d7e6cdbf..ab037bb264 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/mod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/mod.rs @@ -43,7 +43,7 @@ use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::{Context, Ptr}; use pliron::identifier::Identifier; use pliron::linked_list::ContainsLinkedList; -use pliron::location::{Located, Location, Source}; +use pliron::location::{Location, Source}; use pliron::op::Op; use pliron::operation::Operation; use pliron::r#type::{TypeHandle, Typed}; diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/shared.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/shared.rs index 1e4702b076..dc42bf5962 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/shared.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/shared.rs @@ -194,6 +194,98 @@ fn convert_shared_alloc_rejects_conflicting_key_declarations() { ); } +#[test] +fn shared_alloc_alignment_is_at_least_scalar_requirement() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + + append_shared_alloc_typed(&mut ctx, block, "under-aligned-i64", i64_ty, 1, None, 1); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("expected shared global"); + assert_eq!( + global.get_alignment(&ctx), + Some(8), + "typed i64 shared storage must not be weakened below its ABI alignment" + ); +} + +#[test] +fn shared_alloc_alignment_preserves_mir_over_alignment() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let aligned_ty = over_aligned_tuple_ty(&mut ctx); + + append_shared_alloc_typed( + &mut ctx, + block, + "under-aligned-aggregate", + aligned_ty, + 1, + None, + 1, + ); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("expected shared global"); + assert_eq!( + global.get_alignment(&ctx), + Some(32), + "MIR ABI over-alignment must survive shared-storage lowering" + ); +} + +/// The importer hands an opaque descriptor to the dialect as a bare +/// `llvm.array`. No MIR layout describes that leaf, so the LLVM natural +/// layout must supply the storage alignment instead of failing closed. +#[test] +fn shared_alloc_alignment_covers_llvm_array_element() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let descriptor_ty: TypeHandle = ArrayType::get(&ctx, i64_ty, 16).into(); + + append_shared_alloc_typed( + &mut ctx, + block, + "llvm-array-element", + descriptor_ty, + 1, + None, + 1, + ); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .find_map(|op| Operation::get_op::(op, &ctx)) + .expect("expected shared global"); + assert_eq!( + global.get_alignment(&ctx), + Some(8), + "an llvm.array element takes its storage alignment from the LLVM natural layout" + ); +} + #[test] fn shared_alloc_repeated_debug_identity_is_attached_once() { let mut ctx = make_ctx(); @@ -649,11 +741,21 @@ fn append_extern_shared( byte_offset: u64, ) -> Ptr { let i8_ty: TypeHandle = IntegerType::get(ctx, 8, Signedness::Signless).into(); - let shared_i8 = MirPtrType::get_shared(ctx, i8_ty, true); + append_extern_shared_typed(ctx, block, i8_ty, alignment, byte_offset) +} + +fn append_extern_shared_typed( + ctx: &mut Context, + block: Ptr, + pointee_type: TypeHandle, + alignment: u64, + byte_offset: u64, +) -> Ptr { + let shared_ptr = MirPtrType::get_shared(ctx, pointee_type, true); let op = Operation::new( ctx, mir::MirExternSharedOp::get_concrete_op_info(), - vec![shared_i8.into()], + vec![shared_ptr.into()], vec![], vec![], 0, @@ -665,6 +767,101 @@ fn append_extern_shared( op } +#[test] +fn extern_shared_alignment_is_at_least_pointee_requirement() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + + append_extern_shared_typed(&mut ctx, block, i64_ty, 1, 0); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .find(|global| { + global + .get_symbol_name(&ctx) + .to_string() + .starts_with("__dynamic_smem_") + }) + .expect("expected dynamic shared-memory global"); + assert_eq!( + global.get_alignment(&ctx), + Some(8), + "typed i64 dynamic shared storage must not be weakened below its ABI alignment" + ); +} + +#[test] +fn extern_shared_alignment_uses_maximum_typed_requirement() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let aligned_ty = over_aligned_tuple_ty(&mut ctx); + + append_extern_shared_typed(&mut ctx, block, i64_ty, 1, 0); + append_extern_shared_typed(&mut ctx, block, aligned_ty, 1, 0); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .find(|global| { + global + .get_symbol_name(&ctx) + .to_string() + .starts_with("__dynamic_smem_") + }) + .expect("expected dynamic shared-memory global"); + assert_eq!( + global.get_alignment(&ctx), + Some(32), + "one dynamic pool must satisfy the strongest typed alignment requirement" + ); +} + +/// The dynamic pool counterpart of `shared_alloc_alignment_covers_llvm_array_element`: +/// a pointee the MIR side cannot classify still raises the pool. +#[test] +fn extern_shared_alignment_covers_llvm_array_pointee() { + let mut ctx = make_ctx(); + let (module_ptr, block) = build_kernel(&mut ctx, vec![], vec![]); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let descriptor_ty: TypeHandle = ArrayType::get(&ctx, i64_ty, 16).into(); + + append_extern_shared_typed(&mut ctx, block, descriptor_ty, 1, 0); + append_mir_return(&mut ctx, block, vec![]); + + crate::lower_mir_to_llvm(&mut ctx, module_ptr).expect("lowering failed"); + + let top = module_top_block(&ctx, module_ptr); + let global = top + .deref(&ctx) + .iter(&ctx) + .filter_map(|op| Operation::get_op::(op, &ctx)) + .find(|global| { + global + .get_symbol_name(&ctx) + .to_string() + .starts_with("__dynamic_smem_") + }) + .expect("expected dynamic shared-memory global"); + assert_eq!( + global.get_alignment(&ctx), + Some(8), + "an llvm.array pointee takes its pool alignment from the LLVM natural layout" + ); +} + /// A `mir.extern_shared` with a nonzero byte offset addresses into an extern /// `[0 x i8]` addrspace(3) global whose real size only exists at launch, so /// the lowering cannot promise LLVM's `inbounds` allocated-object contract diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/transfer.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/transfer.rs index 85c5e1acbb..61afa19204 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/transfer.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/transfer.rs @@ -20,7 +20,6 @@ use pliron::identifier::Identifier; use pliron::irbuild::dialect_conversion::{DialectConversionRewriter, OperandsInfo}; use pliron::irbuild::inserter::Inserter; use pliron::irbuild::rewriter::Rewriter; -use pliron::location::Located; use pliron::op::Op; use pliron::operation::Operation; use pliron::result::Result; diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/enum_layout.rs b/cuda-oxide/crates/mir-lower/src/convert/types/enum_layout.rs index 2c8c06db4a..64752ce5ce 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/enum_layout.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/enum_layout.rs @@ -15,7 +15,7 @@ use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::Context; use pliron::r#type::TypeHandle; -use super::layout::{make_enum_filler_type, mir_type_contains_i1}; +use super::layout::{alignment_anchor_type, make_enum_filler_type, mir_type_contains_i1}; use super::pointer_storage::{ PointerOverlapRejection, analyze_pointer_overlap, llvm_type_contains_pointer, }; @@ -717,11 +717,12 @@ pub(crate) fn build_enum_slot_map( if natural_align < required_align { // The byte claims alone can under-align the storage, e.g. when the // only claim is an i8 niche carrier inside a 4-aligned enum. Raise - // the struct's alignment with a zero-length anchor field, the same - // mechanism union storage uses; it occupies no bytes, so every slot - // index simply shifts by one. - let anchor_int = IntegerType::get(ctx, (required_align * 8) as u32, Signedness::Signless); - let anchor: TypeHandle = llvm_types::ArrayType::get(ctx, anchor_int.into(), 0).into(); + // the struct's alignment with a zero-length anchor, the same helper + // union storage uses; it occupies no bytes, so every slot index + // simply shifts by one. The anchor is capped at 16 bytes (see + // `alignment_anchor_type`); Rust alignment above that is stamped on + // the memory operations typed as the enum. + let anchor = alignment_anchor_type(ctx, required_align); llvm_fields.insert(0, anchor); for slot in &mut slot_of_claim { *slot += 1; @@ -2509,4 +2510,187 @@ mod tests { "{error}" ); } + + #[test] + fn enum_slot_map_pointer_niche_with_u128_payload_preserves_align16() { + let mut ctx = make_ctx(); + + let logical = mir_uint(&mut ctx, 8); + let wide = mir_uint(&mut ctx, 128); + let u32_ty = mir_uint(&mut ctx, 32); + let pointer: TypeHandle = MirPtrType::get_generic(&mut ctx, u32_ty, false).into(); + + let payload: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![wide, pointer], + vec![0, 1], + vec![0, 16], + 32, + 16, + ) + .into(); + + let enum_ty: TypeHandle = MirEnumType::get_with_encoding( + &mut ctx, + "MaybeWideAndPointer".into(), + logical, + vec![0, 1], + vec![ + EnumVariant::unit("None".into()), + EnumVariant::new_with_layout("Some".into(), vec![payload], vec![0], vec![32]), + ], + EnumEncoding { + tag_offset: 16, + total_size: 32, + abi_align: 16, + layout_kind: EnumLayoutKind::Niche, + carrier_kind: EnumCarrierKind::Pointer, + carrier_width: 64, + untagged_variant: 1, + variant_inhabited: vec![1, 1], + ..EnumEncoding::default() + }, + ) + .into(); + + let map = build_enum_slot_map(&mut ctx, enum_ty).unwrap(); + + let fields = struct_fields(&ctx, map.llvm_struct_ty); + + // The payload decomposes: the u128 leaf gets its own i128 slot, the + // pointer leaf shares the niche carrier, and the tail is filler. The + // i128 slot carries the required 16-byte alignment, so no zero-length + // anchor is needed. + assert_eq!(fields.len(), 3); + assert_eq!( + fields[0] + .deref(&ctx) + .downcast_ref::() + .map(IntegerType::width), + Some(128), + "the u128 leaf should lower to an i128 slot" + ); + assert!( + fields[1].deref(&ctx).is::(), + "the pointer leaf should share the niche carrier slot" + ); + + assert_eq!( + llvm_type_size_align(&ctx, map.llvm_struct_ty), + Some((32, 16)), + "the decomposed payload must preserve rustc's enum size and alignment" + ); + } + + /// `(element count, element width in bits)` of an array-typed filler. + fn filler_array(ctx: &Context, filler: TypeHandle) -> (u64, u32) { + let filler_ref = filler.deref(ctx); + let array = filler_ref + .downcast_ref::() + .expect("enum filler should lower to an LLVM array"); + let elem_width = array + .elem_type() + .deref(ctx) + .downcast_ref::() + .map(IntegerType::width) + .expect("enum filler elements should be integers"); + (array.size(), elem_width) + } + + #[test] + fn enum_filler_16_bytes_at_align4_uses_four_i32s() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 0, 16, 4); + + assert_eq!( + filler_array(&ctx, filler), + (4, 32), + "i64 would raise the enum alignment to 8; i32 is the widest word that fits" + ); + assert_eq!(llvm_type_size_align(&ctx, filler), Some((16, 4))); + } + + #[test] + fn enum_filler_16_bytes_uses_two_i64s_at_align8() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 0, 16, 8); + + let filler_ref = filler.deref(&ctx); + let array = filler_ref + .downcast_ref::() + .expect("16-byte filler should use an LLVM array"); + + assert_eq!(array.size(), 2); + + assert_eq!( + array + .elem_type() + .deref(&ctx) + .downcast_ref::() + .map(IntegerType::width), + Some(64) + ); + + assert_eq!(llvm_type_size_align(&ctx, filler), Some((16, 8))); + } + + #[test] + fn enum_filler_keeps_bytes_at_an_odd_offset() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 1, 16, 8); + + assert_eq!( + filler_array(&ctx, filler), + (16, 8), + "no integer width divides offset 1, so a wider word would move every later field" + ); + assert_eq!(llvm_type_size_align(&ctx, filler), Some((16, 1))); + } + + #[test] + fn enum_filler_keeps_bytes_in_a_byte_aligned_enum() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 0, 16, 1); + + assert_eq!( + filler_array(&ctx, filler), + (16, 8), + "any wider word would raise the enum alignment above rustc's 1" + ); + assert_eq!(llvm_type_size_align(&ctx, filler), Some((16, 1))); + } + + #[test] + fn enum_filler_keeps_bytes_when_no_width_divides_the_size() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 0, 15, 8); + + assert_eq!(filler_array(&ctx, filler), (15, 8)); + assert_eq!(llvm_type_size_align(&ctx, filler), Some((15, 1))); + } + + #[test] + fn enum_filler_24_bytes_uses_three_i64s_at_align8() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 8, 24, 8); + + assert_eq!(filler_array(&ctx, filler), (3, 64)); + assert_eq!(llvm_type_size_align(&ctx, filler), Some((24, 8))); + } + + #[test] + fn enum_filler_12_bytes_uses_three_i32s_at_align4() { + let mut ctx = make_ctx(); + + let filler = make_enum_filler_type(&mut ctx, 4, 12, 4); + + assert_eq!(filler_array(&ctx, filler), (3, 32)); + assert_eq!(llvm_type_size_align(&ctx, filler), Some((12, 4))); + } } diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/layout.rs b/cuda-oxide/crates/mir-lower/src/convert/types/layout.rs index 54083e2424..30f79c8bea 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/layout.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/layout.rs @@ -24,44 +24,75 @@ pub(super) fn make_padding_type(ctx: &mut Context, size: u64) -> TypeHandle { /// Storage for `size` bytes of enum filler at byte `offset`: bytes the payload /// occupies that no typed slot claims. /// -/// `[N x i8]` is byte-exact but costs one leaf *per byte* everywhere a payload -/// value is built, merged, or taken apart. `Option<&[T]>` is the shape that -/// shows it: the niche carrier claims the pointer, the length is left to an -/// 8-byte filler, and building the value then decomposes that length into -/// eight `i8` `insertvalue`s, every block merge carries eight `phi i8`, and -/// reading it back is a chain of `prmt` byte permutes with the aggregate -/// spilled to `.local` in between. Measured on sm_86, it costs the *pointer* -/// too: it rides through the same byte soup, `InferAddressSpaces` can no -/// longer follow it, and the access lands in the generic window (`ld.v2.b64`) -/// instead of `ld.global.v2.b64`. +/// The filler is a vehicle, not a claim. Nothing reads it field-wise: a +/// payload without a typed slot round-trips through memory as a whole +/// aggregate, and both ends of that round trip are byte-exact. What the +/// filler's type does decide is the leaf count of the enum struct, where a +/// leaf is one scalar field left after every array and nested struct is +/// flattened. /// -/// One integer of the same width is the same bytes in one leaf. It is only -/// legal where it moves nothing: +/// The leaf count matters because of `scalarize_block_args`. pliron's +/// `mem2reg` turns an enum local into a block argument of the whole struct, +/// and that pass splits such an argument into scalar block arguments only +/// while the struct has at most `MAX_LEAVES` (16) leaves: `collect_leaves` +/// refuses the 17th leaf and any array over 16 elements. A `[N x i8]` filler +/// costs one leaf per byte. `Option<&[T]>` leaves its length to an 8-byte +/// filler, and as `[8 x i8]` that length was eight `i8` `insertvalue`s, eight +/// `phi i8` at every merge, and a chain of `prmt` byte permutes with the +/// aggregate spilled to `.local` in between. `Option<([u64; 2], &[T])>` +/// stored as `{ [16 x i8], ptr, i64 }` is 18 leaves, so its `Some`/`None` +/// join stays one aggregate phi; `opt` cannot take that apart, +/// `InferAddressSpaces` cannot follow the pointer through the `extractvalue`, +/// and the access lands in the generic window (`ld.v2.b64`) instead of +/// `ld.global.v2.b64`. Stored as `{ [2 x i64], ptr, i64 }` it is 4 leaves, +/// the join becomes scalar phis, and the pointer keeps its address space. /// -/// - the width is 2, 4 or 8 bytes, so the integer is byte-faithful (a multiple -/// of 8 bits, hence no padding of its own). 16 is deliberately excluded: -/// `i128` is legal LLVM but lowers to register pairs on NVPTX, so widening -/// that far trades one win for another cost and wants its own measurement; -/// - `offset` is a multiple of the width, so LLVM inserts no gap ahead of the -/// field and every later field keeps its byte offset; -/// - the width does not exceed the enum's alignment, or the struct's natural -/// alignment would rise above rustc's. +/// The rule: take the widest width among 8, 4 and 2 bytes that divides both +/// `size` and `offset` and does not exceed the enum's alignment; the filler +/// is that integer when it covers `size` alone, else an array of it. rustc +/// types its own padding by the same rule (`type_padding_filler`: the largest +/// of i64, i32 and i16 whose alignment fits). Each condition keeps the layout +/// fixed: /// -/// Anything else keeps the byte array. The filler is a *vehicle*, not a claim: -/// nothing reads it field-wise, because a payload that has no typed slot -/// round-trips through memory as a whole aggregate, and both ends of that -/// round trip are byte-exact. `build_enum_slot_map`'s own size and alignment -/// assertions — hard errors, not debug checks — backstop all three conditions. +/// - the width divides `size`, so the array covers the same bytes; +/// - `offset` is a multiple of the width, so the natural struct walk inserts +/// no gap ahead of the field and every later field keeps its byte offset; +/// - the width does not exceed the enum's alignment, so the struct's natural +/// alignment never rises above rustc's. +/// +/// `i128` is left out: it would impose 16-byte natural alignment, where +/// `[2 x i64]` carries the same 16 bytes at alignment 8. +/// +/// One integer word may cover payload bytes plus Rust padding. Those padding +/// bits are read as undef, the same assumption the 8-byte rule already makes +/// when one `i64` covers a small field and its tail padding. +/// +/// What stays a byte array: a run no width divides, such as the 23 bytes at +/// offset 1 behind a `bool` carrier in `Option<(bool, [u64; 2])>`, and every +/// filler of a byte-aligned enum. What still exceeds the leaf budget even when +/// typed: a filler above 14 i64 words (112 bytes) beside a pointer and a +/// length, since 14 + 2 is the last count that fits. Both keep an aggregate +/// phi and the generic access that comes with it; the `bool` carrier case +/// also keeps its `.local` spill. The budget itself is the follow-up. +/// `build_enum_slot_map`'s size and alignment checks are hard errors and +/// backstop every case. pub(super) fn make_enum_filler_type( ctx: &mut Context, offset: u64, size: u64, abi_align: u64, ) -> TypeHandle { - if matches!(size, 2 | 4 | 8) && offset.is_multiple_of(size) && size <= abi_align.max(1) { - let width = u32::try_from(size * 8).expect("size is at most 8, so width is at most 64"); - return IntegerType::get(ctx, width, Signedness::Signless).into(); + for width in [8u64, 4, 2] { + if size.is_multiple_of(width) && offset.is_multiple_of(width) && width <= abi_align { + let bits = u32::try_from(width * 8).expect("width is at most 8, so at most 64 bits"); + let int_ty: TypeHandle = IntegerType::get(ctx, bits, Signedness::Signless).into(); + if size == width { + return int_ty; + } + return llvm_types::ArrayType::get(ctx, int_ty, size / width).into(); + } } + make_padding_type(ctx, size) } @@ -359,6 +390,46 @@ pub(crate) fn mir_type_abi_align(ctx: &Context, mir_ty: TypeHandle) -> Option TypeHandle { + debug_assert!( + abi_align.is_power_of_two(), + "alignment anchor needs a power-of-two ABI alignment, got {abi_align}" + ); + let anchor_align = abi_align.clamp(1, MAX_NATURAL_ALIGN); + let bits = + u32::try_from(anchor_align * 8).expect("anchor is at most 16 bytes, so at most 128 bits"); + let anchor_int: TypeHandle = IntegerType::get(ctx, bits, Signedness::Signless).into(); + llvm_types::ArrayType::get(ctx, anchor_int, 0).into() +} + /// LLVM natural-layout `(size, align)` of an exported LLVM type, in bytes. /// /// Mirrors LLVM's default data layout for nvptx64 (scalars align to their @@ -371,8 +442,9 @@ pub(crate) fn llvm_type_size_align(ctx: &Context, ty: TypeHandle) -> Option<(u64 if let Some(int_ty) = ty_ref.downcast_ref::() { let size = (int_ty.width() as u64).div_ceil(8); - // i8 → 1, i16 → 2, i32 → 4, i64 → 8, i128 → 16. - return Some((size, size.next_power_of_two().min(16))); + // i8 → 1, i16 → 2, i32 → 4, i64 → 8, i128 → 16; wider integers stay + // at 16, see `MAX_NATURAL_ALIGN`. + return Some((size, size.next_power_of_two().min(MAX_NATURAL_ALIGN))); } if ty_ref.is::() { return Some((2, 2)); @@ -444,6 +516,28 @@ pub(crate) fn llvm_type_size_align(ctx: &Context, ty: TypeHandle) -> Option<(u64 None } +/// Storage alignment a value of `mir_ty` needs, in bytes. +/// +/// Two layouts answer this. The MIR aggregate carries rustc's ABI alignment, +/// including a `repr(align(N))` raise that the converted LLVM type cannot +/// express. The converted LLVM type `llvm_ty` carries the natural alignment +/// of every leaf, including LLVM-dialect leaves such as the `llvm.array` the +/// importer uses for opaque descriptors, which no MIR layout describes. +/// Storage must satisfy both, so the larger answer wins. `None` only when +/// neither layout knows the type. +pub(crate) fn required_storage_align( + ctx: &Context, + mir_ty: TypeHandle, + llvm_ty: TypeHandle, +) -> Option { + let rust_align = mir_type_abi_align(ctx, mir_ty); + let llvm_align = llvm_type_size_align(ctx, llvm_ty).map(|(_, align)| align); + match (rust_align, llvm_align) { + (Some(rust), Some(llvm)) => Some(rust.max(llvm)), + (rust, llvm) => rust.or(llvm), + } +} + /// The one walk defining natural (non-packed) LLVM struct placement: each /// field starts at the running offset rounded up to its own alignment. /// @@ -545,4 +639,53 @@ mod tests { "the NVPTX data layout does not define a 24-bit vector alignment" ); } + + #[test] + fn required_storage_align_takes_the_larger_of_rust_abi_and_llvm_natural() { + let mut ctx = make_ctx(); + + // An LLVM-dialect leaf, as the importer emits for opaque descriptors: + // no MIR layout describes it, so the LLVM natural layout answers. + let i64 = llvm_int(&mut ctx, 64); + let llvm_array: TypeHandle = llvm_types::ArrayType::get(&ctx, i64, 16).into(); + assert_eq!( + required_storage_align(&ctx, llvm_array, llvm_array), + Some(8) + ); + + // A `repr(align(32))` tuple whose converted LLVM type is 1-aligned: + // the Rust raise wins. + let byte = mir_uint(&mut ctx, 8); + let marker: TypeHandle = MirStructType::get_with_full_layout( + &mut ctx, + "Align32".into(), + vec![], + vec![], + vec![], + vec![], + 0, + 32, + ) + .into(); + let tuple: TypeHandle = MirTupleType::get_with_layout( + &mut ctx, + vec![marker, byte], + vec![0, 1], + vec![0, 0], + 32, + 32, + ) + .into(); + let i8 = llvm_int(&mut ctx, 8); + assert_eq!(required_storage_align(&ctx, tuple, i8), Some(32)); + + // A plain byte array carries no Rust raise; the LLVM element answers. + let bytes: TypeHandle = MirArrayType::get(&mut ctx, byte, 4).into(); + let llvm_bytes: TypeHandle = llvm_types::ArrayType::get(&ctx, i8, 4).into(); + assert_eq!(required_storage_align(&ctx, bytes, llvm_bytes), Some(1)); + + // Neither layout knows a void: fail closed. + let void: TypeHandle = llvm_types::VoidType::get(&ctx).into(); + assert_eq!(required_storage_align(&ctx, byte, void), None); + } } diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs b/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs index 1d112dfa60..af6089f3cc 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs @@ -135,11 +135,16 @@ pub(crate) use global_layout::{ pub(crate) use layout::{ llvm_byte_faithful_twin, llvm_type_contains_i1, llvm_type_is_byte_faithful, llvm_type_size_align, mir_element_stride, mir_type_abi_align, natural_struct_layout, + required_storage_align, }; +// Outside this module the anchor cap is read only by test fixtures; a plain +// re-export would trip `unused_imports` in non-test builds. +#[cfg(test)] +pub(crate) use layout::MAX_NATURAL_ALIGN; pub(crate) use pointer_storage::{ llvm_packed_struct_contains_pointer_in_address_space, llvm_type_contains_pointer_in_address_space, }; pub use struct_layout::struct_value_lowering_is_byte_faithful; pub(crate) use struct_layout::{StructLayoutInfo, StructSlotMap, build_struct_slot_map}; -pub(crate) use union_storage::build_union_storage_type; +pub(crate) use union_storage::{build_union_storage_type, union_field_storage_type}; diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/union_storage.rs b/cuda-oxide/crates/mir-lower/src/convert/types/union_storage.rs index aa821690ff..455769b517 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/union_storage.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/union_storage.rs @@ -7,13 +7,49 @@ use dialect_mir::types::MirUnionType; use llvm_export::types as llvm_types; -use pliron::builtin::types::{IntegerType, Signedness}; use pliron::context::Context; use pliron::r#type::TypeHandle; -use super::layout::make_padding_type; +use super::layout::{alignment_anchor_type, make_padding_type}; use super::pointer_storage::llvm_type_contains_pointer; use super::{convert_type, llvm_type_is_byte_faithful, llvm_type_size_align}; +use crate::convert::target_stable_storage::{StorageRewriteOptions, target_stable_storage_type}; + +/// Maximum number of shared-pointer leaves introduced by fixed-array expansion +/// while rebuilding one union field at its semantic/physical storage boundary. +/// +/// Direct pointers and struct nesting remain proportional to source structure. +/// Arrays can encode arbitrarily many extract/cast/insert sequences compactly, +/// so only array-expanded AS3 leaves count against this code-shape budget. +pub(crate) const MAX_UNION_FIELD_ARRAY_REWRITE_LEAVES: u64 = 16; + +/// Return the target-stable physical type used for one union field. +/// +/// Shared-memory pointers become generic pointers recursively so the union's +/// byte image is independent of whether the eventual backend uses legacy +/// 64-bit AS3 pointers or modern NVVM's 32-bit AS3 representation. +pub(crate) fn union_field_storage_type( + ctx: &mut Context, + semantic_ty: TypeHandle, +) -> Result { + let rewrite = target_stable_storage_type( + ctx, + semantic_ty, + StorageRewriteOptions { + canonicalize_bool: false, + }, + "union field storage", + )?; + + if rewrite.array_shared_pointer_leaves > MAX_UNION_FIELD_ARRAY_REWRITE_LEAVES { + return Err(anyhow::anyhow!( + "union field storage: arrays containing shared-memory pointers are not supported above the bounded rewrite limit; rewrite requires {} pointer conversions, supported bound is {MAX_UNION_FIELD_ARRAY_REWRITE_LEAVES}", + rewrite.array_shared_pointer_leaves + )); + } + + Ok(rewrite.ty) +} /// Build byte-exact LLVM storage for a Rust union. /// @@ -24,8 +60,9 @@ use super::{convert_type, llvm_type_is_byte_faithful, llvm_type_size_align}; /// alignment without consuming storage. Pointer-bearing fields are preferred /// so an ordinary union copy keeps LLVM pointer provenance. /// -/// NVPTX gives scalar integers natural alignments up to 16 bytes. Stronger -/// Rust alignment is carried explicitly on memory operations because LLVM +/// The anchor comes from [`alignment_anchor_type`], which caps it at +/// `MAX_NATURAL_ALIGN` (16 bytes) for the reasons given there. Stronger Rust +/// alignment is carried explicitly on memory operations because LLVM /// aggregate types cannot encode over-alignment; the storage type still keeps /// the union's exact size and therefore its array stride. pub(crate) fn build_union_storage_type( @@ -92,12 +129,22 @@ pub(crate) fn build_union_storage_type( } pointer_carrier = Some(llvm_field_ty); } - fields.push((llvm_field_ty, field_size, field_align, contains_pointer)); + let storage_field_ty = union_field_storage_type(ctx, llvm_field_ty)?; + // The storage rewrite only swaps pointer leaves for pointer leaves, + // and lowering sizes every pointer at 8 bytes whatever its address + // space, so storage and semantic layout always agree. + debug_assert_eq!( + llvm_type_size_align(ctx, storage_field_ty), + Some((field_size, field_align)), + "union `{}` field {} target-stable storage changed its size or alignment", + union_ty.name(), + index + ); + + fields.push((storage_field_ty, field_size, field_align, contains_pointer)); } - let storage_align = align.min(16); - let anchor_int = IntegerType::get(ctx, (storage_align * 8) as u32, Signedness::Signless); - let anchor: TypeHandle = llvm_types::ArrayType::get(ctx, anchor_int.into(), 0).into(); + let anchor = alignment_anchor_type(ctx, align); let mut storage_fields = vec![anchor]; if size > 0 { let representative = fields @@ -273,4 +320,84 @@ mod tests { let llvm_array = convert_type(&mut ctx, array).unwrap(); assert_eq!(llvm_type_size_align(&ctx, llvm_array), Some((96, 16))); } + + #[test] + fn union_field_storage_accepts_shared_pointer_array_at_rewrite_bound() { + let mut ctx = make_ctx(); + let u32_ty = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); + let pointers: TypeHandle = + MirArrayType::get(&mut ctx, shared, MAX_UNION_FIELD_ARRAY_REWRITE_LEAVES).into(); + + let semantic = convert_type(&mut ctx, pointers) + .expect("shared-pointer array must convert to an LLVM semantic type"); + let storage = union_field_storage_type(&mut ctx, semantic) + .expect("the exact bounded-array rewrite limit must remain supported"); + + let storage_ref = storage.deref(&ctx); + let array = storage_ref + .downcast_ref::() + .expect("rewritten field must remain an LLVM array"); + let element_ref = array.elem_type().deref(&ctx); + let pointer = element_ref + .downcast_ref::() + .expect("rewritten array elements must remain pointer-typed"); + + assert_eq!( + pointer.address_space(), + llvm_types::address_space::GENERIC, + "bounded AS3 array elements must use generic pointer storage" + ); + } + + #[test] + fn union_field_storage_rejects_shared_pointer_array_above_rewrite_bound() { + let mut ctx = make_ctx(); + let u32_ty = mir_uint(&mut ctx, 32); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); + let pointers: TypeHandle = + MirArrayType::get(&mut ctx, shared, MAX_UNION_FIELD_ARRAY_REWRITE_LEAVES + 1).into(); + + let semantic = convert_type(&mut ctx, pointers) + .expect("shared-pointer array must convert to an LLVM semantic type"); + let error = union_field_storage_type(&mut ctx, semantic) + .expect_err("array-expanded AS3 leaves above the budget must fail closed"); + + assert!( + error.to_string().contains("bounded rewrite limit"), + "{error}" + ); + } + + #[test] + fn union_storage_genericizes_target_dependent_shared_pointer_carrier() { + let mut ctx = make_ctx(); + let u32_ty = mir_uint(&mut ctx, 32); + let u64_ty = mir_uint(&mut ctx, 64); + let shared: TypeHandle = MirPtrType::get_shared(&mut ctx, u32_ty, false).into(); + + let union_ty = MirUnionType::get( + &mut ctx, + "SharedPointerBits".into(), + vec!["ptr".into(), "bits".into()], + vec![shared, u64_ty], + 8, + 8, + ); + let union_data = union_ty.deref(&ctx).clone(); + + let storage = build_union_storage_type(&mut ctx, &union_data) + .expect("direct shared pointer must use target-stable union storage"); + let fields = struct_fields(&ctx, storage); + let carrier_ref = fields[1].deref(&ctx); + let carrier = carrier_ref + .downcast_ref::() + .expect("shared pointer carrier must remain pointer-typed"); + assert_eq!( + carrier.address_space(), + llvm_types::address_space::GENERIC, + "shared-pointer union storage must use a target-stable generic pointer carrier" + ); + assert_eq!(llvm_type_size_align(&ctx, storage), Some((8, 8))); + } } diff --git a/cuda-oxide/crates/mir-lower/src/lib.rs b/cuda-oxide/crates/mir-lower/src/lib.rs index 757e3fc132..82fa281c4b 100644 --- a/cuda-oxide/crates/mir-lower/src/lib.rs +++ b/cuda-oxide/crates/mir-lower/src/lib.rs @@ -143,7 +143,6 @@ use pliron::{ DialectConversion, DialectConversionRewriter, OperandsInfo, apply_dialect_conversion, }, irbuild::{listener::Recorder, rewriter::IRRewriter}, - location::Located, op::{Op, op_cast}, operation::Operation, opts::simplify_cfg::remove_blocks_inside_op, @@ -304,7 +303,7 @@ impl DialectConversion for MirToLlvmConversionDriver { operands_info: &OperandsInfo, ) -> Result<()> { let opid = Operation::get_opid(op, ctx); - let loc = op.deref(ctx).loc(); + let loc = op.deref(ctx).loc().clone(); // Special-case ops that need CUDA pass-level state. if opid == dialect_mir::ops::MirFuncOp::get_opid_static() { diff --git a/cuda-oxide/crates/mir-lower/src/lowering.rs b/cuda-oxide/crates/mir-lower/src/lowering.rs index a098ffab88..b0877b5b7e 100644 --- a/cuda-oxide/crates/mir-lower/src/lowering.rs +++ b/cuda-oxide/crates/mir-lower/src/lowering.rs @@ -30,11 +30,12 @@ use crate::context::{DynamicSmemAlignmentMap, SharedGlobalsMap}; use crate::convert::types::{ StructLayoutInfo, build_struct_slot_map, convert_function_type, convert_type, is_kernel_func, - is_zero_sized_type, llvm_type_size_align, mir_type_abi_align, transparent_scalar_field, + is_zero_sized_type, llvm_type_size_align, mir_type_abi_align, required_storage_align, + transparent_scalar_field, }; use dialect_mir::ops::MirFuncOp; -use dialect_mir::types::{MirDisjointSliceType, MirSliceType, MirStructType}; +use dialect_mir::types::{MirDisjointSliceType, MirPtrType, MirSliceType, MirStructType}; use llvm_export::ops as llvm; use pliron::{ basic_block::BasicBlock, @@ -46,10 +47,11 @@ use pliron::{ rewriter::Rewriter, }, linked_list::ContainsLinkedList, + location::Location, op::Op, operation::Operation, result::Result, - r#type::TypeHandle, + r#type::{TypeHandle, Typed}, value::Value, }; use reserved_oxide_symbols::{ @@ -303,7 +305,7 @@ pub fn convert_func( } propagate_return_abi_alignment(ctx, &llvm_func, return_abi_alignment); - propagate_alwaysinline_attr(ctx, op, &llvm_func); + lower_inline_intent_attr(ctx, &mir_func, &llvm_func); let llvm_entry = llvm_func.get_or_create_entry_block(ctx); @@ -314,7 +316,7 @@ pub fn convert_func( // Pre-scan MIR blocks for max dynamic shared memory alignment. // Must happen BEFORE inline_region empties the MIR region. let mir_blocks: Vec<_> = mir_region.deref(ctx).iter(ctx).collect(); - let body_max_align = compute_max_dynamic_smem_alignment(ctx, &mir_blocks); + let body_max_align = compute_max_dynamic_smem_alignment(ctx, &mir_blocks)?; let contract_min_align = dynamic_shared_alignment_attr(ctx, op); let max_align = match (body_max_align, contract_min_align) { (Some(body), Some(contract)) => Some(body.max(contract)), @@ -419,7 +421,10 @@ fn kernel_reference_param_validities( mir_func: &MirFuncOp, parameter_layout: &[KernelParameterLayout], llvm_func_type: pliron::r#type::TypedHandle, -) -> std::result::Result, anyhow::Error> { +) -> std::result::Result< + Vec<(usize, dialect_mir::attributes::ReferenceParamValidityAttr)>, + anyhow::Error, +> { use pliron::builtin::type_interfaces::FunctionTypeInterface; let llvm_args = llvm_func_type.deref(ctx).arg_types().to_vec(); @@ -442,7 +447,7 @@ fn kernel_reference_param_validities( "kernel source argument {source_index} carries reference validity but its first ABI component is not an LLVM pointer" )); } - result.push((llvm_index, validity.0)); + result.push((llvm_index, validity)); } Ok(result) } @@ -450,14 +455,18 @@ fn kernel_reference_param_validities( fn propagate_kernel_reference_param_validities( ctx: &mut Context, llvm_func: &llvm::FuncOp, - validities: &[(usize, u64)], + validities: &[(usize, dialect_mir::attributes::ReferenceParamValidityAttr)], ) { - for &(index, alignment) in validities { + for &(index, validity) in validities { llvm::set_kernel_reference_param_validity( ctx, llvm_func.get_operation(), index, - llvm::KernelReferenceParamValidityAttr(alignment), + llvm::KernelReferenceParamValidityAttr { + pointee_alignment: validity.pointee_alignment, + no_alias: validity.no_alias, + dereferenceable_bytes: validity.dereferenceable_bytes, + }, ); } } @@ -867,33 +876,34 @@ fn propagate_return_abi_alignment( .set(key, IntegerAttr::new(u64_ty, value)); } -/// Propagate the `alwaysinline` attribute from MIR func to LLVM func. +/// Lower the MIR func's `#[inline]` request (`MirFuncOp::inline_intent`) to +/// its LLVM function attribute in the LLVM func's `llvm_func_attrs`, which +/// `llvm-export` renders like any other. The mapping is rustc's own: +/// `#[inline]` is `inlinehint`, `#[inline(always)]` is `alwaysinline` and +/// `#[inline(never)]` is `noinline`. This is the only place that names the +/// LLVM keywords; the dialect records the request, not the LLVM attribute. /// -/// Set on the MIR func op by `mir-importer` when the source Rust function -/// carries `#[inline(always)]`. The LLVM exporter then emits the -/// `alwaysinline` keyword on the `define` line. Existing `opt -O2` runs can -/// honor that attribute before `llc`, but this propagation is not a mandatory -/// always-inline pass. The goal is to preserve Rust's inline intent for device -/// helpers rather than leaving helper boundaries solely to optimizer +/// `opt -O2` honors the keyword before `llc`; this adds no inliner pass. The +/// goal is to preserve Rust's inline intent for device helpers in both +/// directions, rather than leaving helper boundaries solely to optimizer /// heuristics. -fn propagate_alwaysinline_attr( - ctx: &mut Context, - mir_op: Ptr, - llvm_func: &llvm::FuncOp, -) { - let key: pliron::identifier::Identifier = "alwaysinline".try_into().unwrap(); - let attr_opt = mir_op - .deref(ctx) - .attributes - .get::(&key) - .cloned(); - if let Some(attr) = attr_opt { - llvm_func - .get_operation() - .deref_mut(ctx) - .attributes - .set(key, attr); - } +fn lower_inline_intent_attr(ctx: &mut Context, mir_func: &MirFuncOp, llvm_func: &llvm::FuncOp) { + use dialect_mir::attributes::InlineIntentAttr; + + let Some(intent) = mir_func.inline_intent(ctx) else { + return; + }; + let keyword = match intent { + InlineIntentAttr::Hint => "inlinehint", + InlineIntentAttr::Always => "alwaysinline", + InlineIntentAttr::Never => "noinline", + }; + let mut attrs = llvm_func + .get_attr_llvm_func_attrs(ctx) + .map(|attrs| (*attrs).clone()) + .unwrap_or_default(); + attrs.set(keyword, llvm::LlvmAttrValue::Unit); + llvm_func.set_attr_llvm_func_attrs(ctx, attrs); } // ===================================================================== @@ -1260,32 +1270,67 @@ fn insert_op_sequentially( /// Compute the maximum dynamic shared memory alignment across all /// `MirExternSharedOp` operations in a function. /// +/// Each reference is first raised to the storage alignment of its pointee +/// type, so the one pool satisfies its strongest typed use; the caller +/// folds in the launch contract. The verifier has already rejected a +/// requested alignment that is not a non-zero power of two. +/// /// This pre-pass must run BEFORE `inline_region` moves the blocks, since /// it iterates the MIR blocks directly. The result is stored in /// [`DynamicSmemAlignmentMap`] so that later per-op converters can /// create the global with the correct alignment. fn compute_max_dynamic_smem_alignment( - ctx: &Context, + ctx: &mut Context, mir_blocks: &[Ptr], -) -> Option { - let mut max_alignment: Option = None; - +) -> Result> { + // Collect first: converting a pointee type needs the context mutably, + // which the block iteration cannot hold. + let mut references: Vec<(u64, TypeHandle, Location)> = Vec::new(); for mir_block in mir_blocks { for op in mir_block.deref(ctx).iter(ctx) { let op_id = Operation::get_opid(op, ctx); - if op_id == dialect_mir::ops::MirExternSharedOp::get_opid_static() { - let extern_shared = dialect_mir::ops::MirExternSharedOp::new(op); - let alignment = extern_shared.get_alignment_value(ctx); - - max_alignment = Some(match max_alignment { - Some(current_max) => current_max.max(alignment), - None => alignment, - }); + if op_id != dialect_mir::ops::MirExternSharedOp::get_opid_static() { + continue; } + let extern_shared = dialect_mir::ops::MirExternSharedOp::new(op); + let requested_alignment = extern_shared.get_alignment_value(ctx); + let op_ref = op.deref(ctx); + let result_ty = op_ref.get_result(0).get_type(ctx); + let pointee = result_ty + .deref(ctx) + .downcast_ref::() + .map(|ptr_ty| ptr_ty.pointee) + .ok_or_else(|| { + pliron::create_error!( + op_ref.loc(), + pliron::result::ErrorKind::VerificationFailed, + "mir.extern_shared result is not a MIR pointer" + ) + })?; + references.push((requested_alignment, pointee, op_ref.loc().clone())); } } - max_alignment + let mut max_alignment: Option = None; + for (requested_alignment, pointee, loc) in references { + let llvm_pointee = convert_type(ctx, pointee).map_err(|e| { + pliron::create_error!(&loc, pliron::result::ErrorKind::VerificationFailed, "{e}") + })?; + let required_alignment = + required_storage_align(ctx, pointee, llvm_pointee).ok_or_else(|| { + let pointee_display = pointee.deref(ctx).disp(ctx).to_string(); + pliron::create_error!( + &loc, + pliron::result::ErrorKind::VerificationFailed, + "mir.extern_shared pointee type `{pointee_display}` has no known ABI \ + alignment; refusing to guess the pool alignment" + ) + })?; + let alignment = requested_alignment.max(required_alignment); + max_alignment = Some(max_alignment.map_or(alignment, |current| current.max(alignment))); + } + + Ok(max_alignment) } // ===================================================================== @@ -1599,7 +1644,7 @@ mod reference_param_validity_tests { use super::*; use dialect_mir::{ attributes::ReferenceParamValidityAttr, - types::{MirPointerKind, MirPtrType, MirSliceType}, + types::{MirArrayType, MirPointerKind, MirPtrType, MirSliceType}, }; use pliron::{ builtin::{ @@ -1630,8 +1675,20 @@ mod reference_param_validity_tests { MirPtrType::get_generic_with_kind(&mut ctx, f32_ty, false, MirPointerKind::RawConst) .into(); - let mir_func_type = - FunctionType::get(&ctx, vec![shared_ref, shared_slice, raw_ptr], vec![]); + let zst: TypeHandle = MirArrayType::get(&mut ctx, f32_ty, 0).into(); + let mir_func_type = FunctionType::get( + &ctx, + vec![ + shared_ref, + zst, + shared_slice, + raw_ptr, + zst, + shared_slice, + shared_ref, + ], + vec![], + ); let op = Operation::new( &mut ctx, MirFuncOp::get_concrete_op_info(), @@ -1641,8 +1698,19 @@ mod reference_param_validity_tests { 1, ); let mir_func = MirFuncOp::new(&mut ctx, op, TypeAttr::new(mir_func_type.into())); - mir_func.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr(4)); - mir_func.set_reference_param_validity(&mut ctx, 1, ReferenceParamValidityAttr(4)); + let sized = ReferenceParamValidityAttr { + pointee_alignment: 4, + no_alias: true, + dereferenceable_bytes: 4, + }; + let slice = ReferenceParamValidityAttr { + dereferenceable_bytes: 0, + ..sized + }; + let aliasable = ReferenceParamValidityAttr::new(4); + for (source_index, validity) in [(0, sized), (2, slice), (5, aliasable), (6, sized)] { + mir_func.set_reference_param_validity(&mut ctx, source_index, validity); + } let llvm_func_type = convert_function_type(&mut ctx, mir_func_type, true).expect("kernel type lowers"); @@ -1651,14 +1719,52 @@ mod reference_param_validity_tests { kernel_reference_param_validities(&mut ctx, &mir_func, &layout, llvm_func_type) .expect("reference validity mapping succeeds"); - assert_eq!(mapped, vec![(0, 4), (1, 4)]); + assert_eq!( + mapped, + vec![(0, sized), (1, slice), (4, aliasable), (6, sized)] + ); + let llvm_func = + llvm::FuncOp::new(&mut ctx, "transport".try_into().unwrap(), llvm_func_type); + propagate_kernel_reference_param_validities(&mut ctx, &llvm_func, &mapped); + let entries = + llvm::kernel_reference_param_validity_entries(&ctx, llvm_func.get_operation()).unwrap(); + assert_eq!( + entries, + vec![ + ( + 0, + llvm::KernelReferenceParamValidityAttr { + pointee_alignment: 4, + no_alias: true, + dereferenceable_bytes: 4, + } + ), + ( + 1, + llvm::KernelReferenceParamValidityAttr { + pointee_alignment: 4, + no_alias: true, + dereferenceable_bytes: 0, + } + ), + (4, llvm::KernelReferenceParamValidityAttr::new(4)), + ( + 6, + llvm::KernelReferenceParamValidityAttr { + pointee_alignment: 4, + no_alias: true, + dereferenceable_bytes: 4, + } + ), + ] + ); use pliron::builtin::type_interfaces::FunctionTypeInterface; let llvm_args = llvm_func_type.deref(&ctx).arg_types().to_vec(); assert_eq!( llvm_args.len(), - 4, - "slice lowers to data pointer + length while raw pointer stays one parameter" + 7, + "slices expand, ZSTs disappear, and raw pointers stay bare" ); assert!( llvm_args[0] @@ -1896,9 +2002,9 @@ mod grid_constant_abi_tests { 1, ); let function = MirFuncOp::new(&mut ctx, op, TypeAttr::new(function_type.into())); - function.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr(1)); - function.set_reference_param_validity(&mut ctx, 2, ReferenceParamValidityAttr(64)); - function.set_reference_param_validity(&mut ctx, 4, ReferenceParamValidityAttr(64)); + function.set_reference_param_validity(&mut ctx, 0, ReferenceParamValidityAttr::new(1)); + function.set_reference_param_validity(&mut ctx, 2, ReferenceParamValidityAttr::new(64)); + function.set_reference_param_validity(&mut ctx, 4, ReferenceParamValidityAttr::new(64)); record(&mut ctx, op, 2, payload, 64); record(&mut ctx, op, 4, payload, 64); let llvm_type = convert_function_type(&mut ctx, function_type, true).unwrap(); @@ -1920,7 +2026,14 @@ mod grid_constant_abi_tests { assert_eq!(llvm_type_size_align(&ctx, grid[0].pointee).unwrap().0, 128); let validities = kernel_reference_param_validities(&mut ctx, &function, &layout, llvm_type).unwrap(); - assert_eq!(validities, vec![(0, 1), (2, 64), (4, 64)]); + assert_eq!( + validities, + vec![ + (0, ReferenceParamValidityAttr::new(1)), + (2, ReferenceParamValidityAttr::new(64)), + (4, ReferenceParamValidityAttr::new(64)), + ] + ); // A stale pointee annotation cannot silently change the launch packet size. record(&mut ctx, op, 2, byte, 64); diff --git a/cuda-oxide/crates/mir-lower/src/scalarize_block_args.rs b/cuda-oxide/crates/mir-lower/src/scalarize_block_args.rs index 908bb1e0db..c26c80d59b 100644 --- a/cuda-oxide/crates/mir-lower/src/scalarize_block_args.rs +++ b/cuda-oxide/crates/mir-lower/src/scalarize_block_args.rs @@ -653,4 +653,79 @@ mod tests { let branch = op_cast::(term_obj.as_ref()).unwrap(); assert_eq!(branch.successor_operands(&ctx, 0).len(), 1); } + + /// Enum storage `{ [count x iN], i64, i64 }`: one filler array beside a + /// pointer-sized carrier and a length (an `i64` stands in for the `ptr` + /// leaf; only the leaf count matters here). Builds a merge of that + /// aggregate fed by two predecessors, runs the pass, and returns how many + /// arguments the merge block carries afterwards: one when the aggregate + /// stays whole, one per leaf when it splits. + fn merge_arguments_for_filler_storage(elem_bits: u32, count: u64) -> usize { + let mut ctx = make_ctx(); + let elem: TypeHandle = IntegerType::get(&ctx, elem_bits, Signedness::Signless).into(); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let filler: TypeHandle = llvm_types::ArrayType::get(&ctx, elem, count).into(); + let agg_ty: TypeHandle = llvm_types::StructType::get_unnamed( + &ctx, + ( + vec![filler, i64_ty, i64_ty], + llvm_types::StructLayout::Unpacked, + ), + ) + .into(); + + let (module_ptr, entry) = build_kernel(&mut ctx, vec![], vec![]); + + let undef_op = llvm::UndefOp::new(&mut ctx, agg_ty); + undef_op.get_operation().insert_at_back(entry, &ctx); + let aggregate = undef_op.get_operation().deref(&ctx).get_result(0); + + let other = append_block(&mut ctx, entry, vec![]); + let merge = append_block(&mut ctx, entry, vec![agg_ty]); + let br_op = llvm::BrOp::new(&mut ctx, merge, vec![aggregate]); + br_op.get_operation().insert_at_back(entry, &ctx); + + let other_undef = llvm::UndefOp::new(&mut ctx, agg_ty); + other_undef.get_operation().insert_at_back(other, &ctx); + let other_agg = other_undef.get_operation().deref(&ctx).get_result(0); + let other_br = llvm::BrOp::new(&mut ctx, merge, vec![other_agg]); + other_br.get_operation().insert_at_back(other, &ctx); + + let return_op = llvm::ReturnOp::new(&mut ctx, None); + return_op.get_operation().insert_at_back(merge, &ctx); + + scalarize_aggregate_block_args(&mut ctx, module_ptr).unwrap(); + + merge.deref(&ctx).get_num_arguments() + } + + /// The filler's type decides whether enum storage fits the leaf budget. + /// `Option<([u64; 2], &[T])>` stored as `{ [16 x i8], ptr, i64 }` is 18 + /// leaves and stays one aggregate argument, which is how the pointer + /// inside it lost its address space; the same bytes typed as `[2 x i64]` + /// are 4 leaves and split. `[14 x i64]` beside two words is the last + /// filler that fits (16 leaves); `[15 x i64]` is one over. + #[test] + fn typed_fillers_decide_whether_enum_storage_splits() { + assert_eq!( + merge_arguments_for_filler_storage(8, 16), + 1, + "{{ [16 x i8], i64, i64 }} is 18 leaves and must stay whole" + ); + assert_eq!( + merge_arguments_for_filler_storage(64, 2), + 4, + "{{ [2 x i64], i64, i64 }} is 4 leaves and must split" + ); + assert_eq!( + merge_arguments_for_filler_storage(64, 14), + 16, + "{{ [14 x i64], i64, i64 }} is 16 leaves, the last count that fits" + ); + assert_eq!( + merge_arguments_for_filler_storage(64, 15), + 1, + "{{ [15 x i64], i64, i64 }} is 17 leaves and must stay whole" + ); + } } diff --git a/cuda-oxide/crates/mir-lower/src/wgmma_deferred_accumulator.rs b/cuda-oxide/crates/mir-lower/src/wgmma_deferred_accumulator.rs index 7c355a5482..80a9d68bbd 100644 --- a/cuda-oxide/crates/mir-lower/src/wgmma_deferred_accumulator.rs +++ b/cuda-oxide/crates/mir-lower/src/wgmma_deferred_accumulator.rs @@ -89,7 +89,6 @@ use pliron::{ rewriter::{IRRewriter, Rewriter}, }, linked_list::ContainsLinkedList, - location::Located, op::{Op, op_cast}, operation::Operation, opts::simplify_cfg::remove_blocks_inside_op, @@ -1309,7 +1308,7 @@ fn apply_value_plan( row_type: TypeHandle, element_type: TypeHandle, ) { - let loc = fence.deref(ctx).loc(); + let loc = fence.deref(ctx).loc().clone(); let rows = linear_accumulator_rows(kind); let accumulator_count = rows * ACCUMULATOR_COLUMNS; debug_assert!( @@ -1411,7 +1410,7 @@ fn apply_pipelined_counted_loop_plan(ctx: &mut Context, plan: PipelinedCountedLo ); debug_assert_eq!(exit.deref(ctx).get_num_arguments(), 0); - let loc = fence.deref(ctx).loc(); + let loc = fence.deref(ctx).loc().clone(); let mut element_pointers = Vec::with_capacity(slot_count); let mut all_accumulator_values = Vec::with_capacity(slot_count * ACCUMULATOR_LEN); for slot in 0..slot_count { @@ -1501,7 +1500,7 @@ fn apply_counted_loop_plan(ctx: &mut Context, plan: CountedLoopPlan) { ); debug_assert_eq!(exit.deref(ctx).get_num_arguments(), 0); - let loc = fence.deref(ctx).loc(); + let loc = fence.deref(ctx).loc().clone(); let (element_pointers, accumulator_values) = load_accumulator_values_before( ctx, preheader_terminator, @@ -1570,7 +1569,7 @@ fn apply_pointer_fallback( descriptors: Vec, ) { let group = WgmmaMmaGroupM64N64K16F32Bf16Op::build(ctx, accumulator, descriptors); - group.deref_mut(ctx).set_loc(fence.deref(ctx).loc()); + group.deref_mut(ctx).set_loc(fence.deref(ctx).loc().clone()); group.insert_before(ctx, wait); erase_original_sequence(ctx, fence, mmas, commit, wait); } @@ -2124,7 +2123,7 @@ fn apply_pipeline_plan(ctx: &mut Context, plan: PipelinePlan) { let final_wait = *waits .last() .expect("pipeline plan has a final wait_group<0>"); - let loc = fence.deref(ctx).loc(); + let loc = fence.deref(ctx).loc().clone(); let mut slot_types = Vec::with_capacity(accumulators.len()); let mut all_accumulator_values = Vec::with_capacity(accumulators.len() * ACCUMULATOR_LEN); diff --git a/cuda-oxide/crates/mir-lower/tests/lowering_test/calls_and_values.rs b/cuda-oxide/crates/mir-lower/tests/lowering_test/calls_and_values.rs index 7840c12933..a22b65f552 100644 --- a/cuda-oxide/crates/mir-lower/tests/lowering_test/calls_and_values.rs +++ b/cuda-oxide/crates/mir-lower/tests/lowering_test/calls_and_values.rs @@ -9,7 +9,7 @@ use pliron::builtin::op_interfaces::{CallOpCallable, CallOpInterface, SymbolOpIn use pliron::builtin::ops::ModuleOp; use pliron::context::Context; use pliron::linked_list::ContainsLinkedList; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; @@ -182,7 +182,7 @@ fn addrspace_coercion_inserts_addrspacecast_at_call_site() -> Result<(), anyhow: found_call = true; assert_eq!( call.get_operation().deref(&ctx).loc(), - expected_call_location + &expected_call_location ); } } diff --git a/cuda-oxide/crates/mir-lower/tests/lowering_test/inline_ptx.rs b/cuda-oxide/crates/mir-lower/tests/lowering_test/inline_ptx.rs index 4adc9e9fbc..c24c217ccb 100644 --- a/cuda-oxide/crates/mir-lower/tests/lowering_test/inline_ptx.rs +++ b/cuda-oxide/crates/mir-lower/tests/lowering_test/inline_ptx.rs @@ -7,7 +7,7 @@ use dialect_nvvm::ops as nvvm; use llvm_export::ops as llvm; use pliron::builtin::op_interfaces::SymbolOpInterface; use pliron::linked_list::ContainsLinkedList; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; @@ -194,7 +194,7 @@ fn test_multi_result_inline_ptx_lowers_to_struct_asm_and_extractvalues() -> Resu if let Some(inline_asm) = Operation::get_op::(op, &ctx) { assert_eq!( inline_asm.get_operation().deref(&ctx).loc(), - expected_location + &expected_location ); assert_eq!( inline_asm @@ -230,7 +230,10 @@ fn test_multi_result_inline_ptx_lowers_to_struct_asm_and_extractvalues() -> Resu } asm_result = Some(result); } else if let Some(extract) = Operation::get_op::(op, &ctx) { - assert_eq!(extract.get_operation().deref(&ctx).loc(), expected_location); + assert_eq!( + extract.get_operation().deref(&ctx).loc(), + &expected_location + ); let aggregate = extract.get_operation().deref(&ctx).get_operand(0); assert_eq!( Some(aggregate), diff --git a/cuda-oxide/crates/mir-lower/tests/lowering_test/sregs_and_warp.rs b/cuda-oxide/crates/mir-lower/tests/lowering_test/sregs_and_warp.rs index 8e06160f3a..dda98b4490 100644 --- a/cuda-oxide/crates/mir-lower/tests/lowering_test/sregs_and_warp.rs +++ b/cuda-oxide/crates/mir-lower/tests/lowering_test/sregs_and_warp.rs @@ -12,7 +12,7 @@ use pliron::builtin::op_interfaces::{ use pliron::builtin::ops::ModuleOp; use pliron::context::Context; use pliron::linked_list::ContainsLinkedList; -use pliron::location::{Located, Location}; +use pliron::location::Location; use pliron::op::Op; use pliron::operation::Operation; @@ -196,7 +196,7 @@ fn test_intrinsic_insertion() -> Result<(), anyhow::Error> { if symbol.to_string() == "llvm_nvvm_read_ptx_sreg_tid_x" ) { found_intrinsic_call = true; - assert_eq!(call.get_operation().deref(&ctx).loc(), expected_location); + assert_eq!(call.get_operation().deref(&ctx).loc(), &expected_location); } } } diff --git a/cuda-oxide/crates/mir-lower/tests/lowering_test/tma.rs b/cuda-oxide/crates/mir-lower/tests/lowering_test/tma.rs index 7cdf3f6d81..1215c19cdb 100644 --- a/cuda-oxide/crates/mir-lower/tests/lowering_test/tma.rs +++ b/cuda-oxide/crates/mir-lower/tests/lowering_test/tma.rs @@ -171,3 +171,147 @@ fn tma_modern_libnvvm_preserves_existing_cluster_pointer_semantics() -> Result<( assert!(!ir.contains("cvta.to.shared::cta"), "{ir}"); Ok(()) } + +/// Lowers the five CTA-local G2S forms with the given destination and +/// barrier address spaces. +fn lower_g2s_cta_forms( + destination_space: u32, + barrier_space: u32, + backend: mir_lower::IntrinsicBackend, +) -> Result<(Context, Ptr), anyhow::Error> { + let mut ctx = make_test_ctx(); + let byte = IntegerType::get(&ctx, 8, Signedness::Signless).into(); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless).into(); + let i64_ty = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let destination = MirPtrType::get(&mut ctx, byte, true, destination_space).into(); + let barrier = MirPtrType::get(&mut ctx, byte, true, barrier_space).into(); + let pointer = MirPtrType::get_generic(&mut ctx, byte, true).into(); + let (module, entry) = build_test_kernel( + &mut ctx, + vec![ + destination, + barrier, + pointer, + i32_ty, + i32_ty, + i32_ty, + i32_ty, + i32_ty, + i64_ty, + ], + ); + let args: Vec<_> = (0..9) + .map(|index| entry.deref(&ctx).get_argument(index)) + .collect(); + for (info, dimensions) in [ + ( + nvvm::CpAsyncBulkTensorG2sCtaTile1dOp::get_concrete_op_info(), + 1, + ), + ( + nvvm::CpAsyncBulkTensorG2sCtaTile2dOp::get_concrete_op_info(), + 2, + ), + ( + nvvm::CpAsyncBulkTensorG2sCtaTile3dOp::get_concrete_op_info(), + 3, + ), + ( + nvvm::CpAsyncBulkTensorG2sCtaTile4dOp::get_concrete_op_info(), + 4, + ), + ( + nvvm::CpAsyncBulkTensorG2sCtaTile5dOp::get_concrete_op_info(), + 5, + ), + ] { + let mut operands = args[..3 + dimensions].to_vec(); + operands.push(args[8]); + Operation::new(&mut ctx, info, vec![], operands, vec![], 0).insert_at_back(entry, &ctx); + } + append_return(&mut ctx, entry); + mir_lower::lower_mir_to_llvm_with_options( + &mut ctx, + module, + mir_lower::LoweringOptions { + intrinsic_backend: backend, + ..Default::default() + }, + ) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok((ctx, module)) +} + +#[test] +fn tma_cta_llvm_nvptx_uses_typed_shared_cta_intrinsics() -> Result<(), anyhow::Error> { + for (destination_space, barrier_space) in [(0, 0), (3, 3), (0, 3), (3, 0)] { + let (ctx, module) = lower_g2s_cta_forms( + destination_space, + barrier_space, + mir_lower::IntrinsicBackend::LlvmNvptx, + )?; + let module = Operation::get_op::(module, &ctx).unwrap(); + let ir = llvm_export::export::export_module_to_string(&ctx, &module).unwrap(); + for dimensions in 1..=5 { + let call = format!( + "call void @llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.{dimensions}d(ptr addrspace(3)" + ); + assert_eq!(ir.matches(&call).count(), 1, "{ir}"); + } + assert!(!ir.contains("addrspace(7)"), "{ir}"); + assert!(!ir.contains("g2s.tile."), "{ir}"); + assert!(!ir.contains("asm sideeffect"), "{ir}"); + } + Ok(()) +} + +#[test] +fn tma_cta_libnvvm_emits_exact_shared_cta_ptx() -> Result<(), anyhow::Error> { + for (destination_space, barrier_space) in [(0, 0), (3, 3)] { + let (ctx, module) = lower_g2s_cta_forms( + destination_space, + barrier_space, + mir_lower::IntrinsicBackend::LibNvvm, + )?; + let module = Operation::get_op::(module, &ctx).unwrap(); + for dialect in [NvvmIrDialect::LegacyLlvm7, NvvmIrDialect::Modern] { + let ir = + export_module_to_string_with_config(&ctx, &module, &NvvmExportConfig::new(dialect)) + .unwrap(); + for dimensions in 1..=5 { + let instruction = format!( + "cp.async.bulk.tensor.{dimensions}d.shared::cta.global.tile.mbarrier::complete_tx::bytes [$0], [$2," + ); + assert_eq!(ir.matches(&instruction).count(), 1, "{ir}"); + } + assert!(!ir.contains("addrspace(7)"), "{ir}"); + assert!(!ir.contains("shared::cluster"), "{ir}"); + } + } + Ok(()) +} + +#[test] +fn tma_cta_rejects_cluster_shared_operands() { + for backend in [ + mir_lower::IntrinsicBackend::LlvmNvptx, + mir_lower::IntrinsicBackend::LibNvvm, + ] { + for (destination_space, barrier_space, role) in [ + (7, 3, "destination"), + (3, 7, "barrier"), + (7, 0, "destination"), + ] { + let error = lower_g2s_cta_forms(destination_space, barrier_space, backend) + .err() + .unwrap_or_else(|| panic!("an addrspace(7) {role} lowered with {backend:?}")) + .to_string(); + assert!( + error.contains("g2s_cta requires a CTA-local shared") + && error.contains(role) + && error.contains("addrspace 7"), + "{error}" + ); + } + } +} diff --git a/cuda-oxide/crates/mir-transforms/src/analyses/induction.rs b/cuda-oxide/crates/mir-transforms/src/analyses/induction.rs index 5b6ae4eef3..719e841d46 100644 --- a/cuda-oxide/crates/mir-transforms/src/analyses/induction.rs +++ b/cuda-oxide/crates/mir-transforms/src/analyses/induction.rs @@ -80,7 +80,7 @@ pub enum CmpPred { impl CmpPred { /// The opposite test: the one that is true exactly when this one is false. /// (`<` becomes `>=`, and so on.) - fn negate(self) -> CmpPred { + pub(crate) fn negate(self) -> CmpPred { match self { CmpPred::Lt => CmpPred::Ge, CmpPred::Le => CmpPred::Gt, @@ -191,7 +191,7 @@ pub(crate) fn edge_operands( /// value underneath and whether the number peeled off was odd (so `!!x` reports /// "not negated" and `!x` reports "negated"). Used so a guard written `!(i < n)` /// is understood the same as `i >= n`. -fn unwrap_not(ctx: &Context, mut v: Value) -> (Value, bool) { +pub(crate) fn unwrap_not(ctx: &Context, mut v: Value) -> (Value, bool) { let mut negated = false; while let Some(def) = v.defining_op() { if Operation::get_op::(def, ctx).is_some() { @@ -206,7 +206,7 @@ fn unwrap_not(ctx: &Context, mut v: Value) -> (Value, bool) { /// If `op` is a comparison (`<`, `<=`, `>`, `>=`), return which one it is and /// its two sides (left, right); `None` for any other op. -fn match_cmp(ctx: &Context, op: Ptr) -> Option<(CmpPred, Value, Value)> { +pub(crate) fn match_cmp(ctx: &Context, op: Ptr) -> Option<(CmpPred, Value, Value)> { let pred = if Operation::get_op::(op, ctx).is_some() { CmpPred::Lt } else if Operation::get_op::(op, ctx).is_some() { diff --git a/cuda-oxide/crates/mir-transforms/src/canonicalize.rs b/cuda-oxide/crates/mir-transforms/src/canonicalize.rs index 771c13b31d..4c771e6a44 100644 --- a/cuda-oxide/crates/mir-transforms/src/canonicalize.rs +++ b/cuda-oxide/crates/mir-transforms/src/canonicalize.rs @@ -24,12 +24,16 @@ //! recompute dominance, loop structure, and induction facts after this rewrite. //! //! A range `for` loop needs one more rewrite, [`thread_iterator_exit_test`]: its -//! exit test is a match on the `Option` from `Iterator::next`. +//! exit test is a match on the `Option` from `Iterator::next`. A signed counter +//! first needs [`fold_checked_step`]: core steps it through `checked_add`. An +//! `isize` step casts through `i64` and back; the importer drops that cast +//! because both are the same dialect type, so the block matches `i64`. use dialect_mir::attributes::MirCastKindAttr; use dialect_mir::ops::{ - MirCastOp, MirCondBranchOp, MirConstantOp, MirConstructEnumOp, MirEnumPayloadOp, MirEqOp, - MirGetDiscriminantOp, MirGotoOp, MirNeOp, MirStorageDeadOp, MirStorageLiveOp, + MirAddOp, MirCastOp, MirCheckedAddOp, MirCheckedSubOp, MirCondBranchOp, MirConstantOp, + MirConstructEnumOp, MirEnumPayloadOp, MirEqOp, MirFieldAddrOp, MirGetDiscriminantOp, MirGotoOp, + MirLoadOp, MirNeOp, MirNotOp, MirStorageDeadOp, MirStorageLiveOp, MirStoreOp, MirSubOp, }; use dialect_mir::types::MirEnumType; use pliron::basic_block::BasicBlock; @@ -43,7 +47,8 @@ use pliron::r#type::{Typed, TypedHandle}; use pliron::value::Value; use rustc_hash::{FxHashMap, FxHashSet}; -use crate::analyses::loop_info::{LoopId, LoopInfo}; +use crate::analyses::induction::{CmpPred, const_i128, match_cmp, unwrap_not}; +use crate::analyses::loop_info::{Loop, LoopId, LoopInfo}; /// Result of trying to put one loop into the single-latch form the unroller /// consumes. @@ -271,7 +276,7 @@ pub(crate) fn merge_backedges( } /// The reason reported for a loop whose header does not exit and that -/// [`thread_iterator_exit_test`] does not recognize. +/// [`thread_iterator_exit_test`] does not recognize as a range `for`. fn unrecognized_iterator_loop() -> CanonicalizeOutcome { CanonicalizeOutcome::Unsupported( "its exit test is not in the loop header, and it is not a range `for` loop over `a..b`" @@ -279,6 +284,23 @@ fn unrecognized_iterator_loop() -> CanonicalizeOutcome { ) } +/// A range `for` loop whose step is not the foldable step of 1. This is still +/// a range loop: the warning must not describe it as something else. +fn unrecognized_range_step() -> CanonicalizeOutcome { + CanonicalizeOutcome::Unsupported( + "its exit test is not in the loop header, and its range `for` step was not recognized" + .into(), + ) +} + +/// Whether `arm` holds the checked add or sub of a signed range step. +fn arm_has_checked_step(ctx: &Context, arm: Ptr) -> bool { + arm.deref(ctx).iter(ctx).any(|op| { + Operation::get_op::(op, ctx).is_some() + || Operation::get_op::(op, ctx).is_some() + }) +} + /// Bound on the blocks followed through one discriminant match, so a cycle of /// match-like blocks cannot hang the pass. const MAX_MATCH_BLOCKS: usize = 16; @@ -333,6 +355,8 @@ struct IteratorArm { /// Returns `Unchanged` when the header exits, and `Unsupported`, with the IR /// untouched, when the header does not exit and the loop misses this shape or /// a value from the match is read where both the body and the exit reach. +/// `Changed` also covers [`fold_checked_step`], which leaves the match for the +/// next round. pub(crate) fn thread_iterator_exit_test( ctx: &mut Context, info: &LoopInfo, @@ -353,6 +377,20 @@ pub(crate) fn thread_iterator_exit_test( return CanonicalizeOutcome::Unchanged; } + // A signed counter's arm ends in a branch on an overflow flag, not in + // `goto join(..)`. Fold that into a plain step; the caller's CFG cleanup + // merges the arm with the block behind the branch, and the next round sees + // the shape above. + let steps: Vec = (0..2) + .filter_map(|index| match_checked_step(ctx, lp, header_term, index)) + .collect(); + if !steps.is_empty() { + for step in steps { + fold_checked_step(ctx, step); + } + return CanonicalizeOutcome::Changed; + } + // Each arm is entered only from the header and ends in `goto join(..)`. let mut join = None; for &arm in &arm_blocks { @@ -363,6 +401,15 @@ pub(crate) fn thread_iterator_exit_test( return unrecognized_iterator_loop(); }; if Operation::get_op::(term, ctx).is_none() { + // A signed range step ends in a branch on its overflow flag. If + // that step was not folded, the loop is still a range `for`; only + // the step is unrecognized. + if arm_blocks + .iter() + .any(|candidate| arm_has_checked_step(ctx, *candidate)) + { + return unrecognized_range_step(); + } return unrecognized_iterator_loop(); } let target = term.deref(ctx).get_successor(0); @@ -590,6 +637,228 @@ pub(crate) fn thread_iterator_exit_test( CanonicalizeOutcome::Changed } +/// A header successor that steps a signed counter the way core's +/// `Step::forward_unchecked` does: see [`match_checked_step`]. +struct CheckedStep { + /// The header successor that holds the step. + arm: Ptr, + /// Its branch on the overflow flag. + branch: Ptr, + /// The block the branch enters when nothing overflowed, and the operands it + /// passes there. + in_range: Ptr, + in_range_operands: Vec, + /// The counter and the amount it steps by, and whether it steps down. + counter: Value, + amount: Value, + subtract: bool, + /// The sum read back from the stack slot the checked tuple went through. + sum: Value, + slot: Value, + /// The operations that build the tuple, store it and read it back, in block + /// order. + chain: Vec>, +} + +/// Find the step of a signed range `for` loop in the header successor +/// `index`. After mem2reg, core's signed `forward_unchecked` (`backward_unchecked` +/// for `.rev()`) leaves this in the arm, with the tuple in a stack slot that +/// mem2reg does not promote: +/// +/// ```text +/// arm: t = checked_add(i, 1); store slot, t +/// x = load field_addr(slot, 0); f = load field_addr(slot, 1) +/// cond_br not(f) [in_range, overflow] +/// in_range: o = Some(i); goto join(o, x) +/// overflow: ... (leaves the loop) +/// ``` +/// +/// The arm must hold nothing else but constants and storage markers, step the +/// counter the header tests by one, and be entered only when the header's test +/// `lo < hi` holds. `lo + 1` then cannot overflow (`lo < hi <= MAX`) and nor +/// can `hi - 1` (`MIN <= lo < hi`), so the overflow edge is dead. Only an edge +/// that leaves the loop is folded; one that stays in it is not `unwrap_unchecked`'s +/// unreachable path. +fn match_checked_step( + ctx: &Context, + lp: &Loop, + header_term: Ptr, + index: usize, +) -> Option { + let arm = header_term.deref(ctx).get_successor(index); + let branch = arm.deref(ctx).get_terminator(ctx)?; + Operation::get_op::(branch, ctx)?; + if arm.uses(ctx).len() != 1 { + return None; + } + + let mut checked = None; + let mut slot = None; + let mut addresses: [Option; 2] = [None; 2]; + let mut loaded: [Option; 2] = [None; 2]; + let mut flag = None; + let mut chain = Vec::new(); + for op in arm.deref(ctx).iter(ctx) { + if op == branch + || Operation::get_op::(op, ctx).is_some() + || Operation::get_op::(op, ctx).is_some() + || Operation::get_op::(op, ctx).is_some() + { + continue; + } + let operation = op.deref(ctx); + if Operation::get_op::(op, ctx).is_some() + || Operation::get_op::(op, ctx).is_some() + { + if checked.is_some() { + return None; + } + let subtract = Operation::get_op::(op, ctx).is_some(); + checked = Some(( + operation.get_result(0), + operation.get_operand(0), + operation.get_operand(1), + subtract, + )); + } else if let Some(store) = Operation::get_op::(op, ctx) { + let (pair, ..) = checked?; + if slot.is_some() || store.is_volatile(ctx) || operation.get_operand(1) != pair { + return None; + } + slot = Some(operation.get_operand(0)); + } else if let Some(address) = Operation::get_op::(op, ctx) { + let field = address.get_attr_field_index(ctx)?.0 as usize; + if slot != Some(operation.get_operand(0)) || field > 1 || addresses[field].is_some() { + return None; + } + addresses[field] = Some(operation.get_result(0)); + } else if let Some(load) = Operation::get_op::(op, ctx) { + let field = addresses + .iter() + .position(|address| *address == Some(operation.get_operand(0)))?; + if load.is_volatile(ctx) || loaded[field].is_some() { + return None; + } + loaded[field] = Some(operation.get_result(0)); + } else if Operation::get_op::(op, ctx).is_some() { + if flag.is_some() || Some(operation.get_operand(0)) != loaded[1] { + return None; + } + flag = Some(operation.get_result(0)); + } else { + return None; + } + chain.push(op); + } + let (_, counter, amount, subtract) = checked?; + let (slot, sum, flag) = (slot?, loaded[0]?, flag?); + + // The branch enters `in_range` when the flag is clear. + let successors: Vec<_> = branch.deref(ctx).successors().collect(); + if branch.deref(ctx).get_operand(0) != flag || successors.len() != 2 { + return None; + } + let (in_range, overflow) = (successors[0], successors[1]); + if in_range == overflow + || !lp.blocks.contains(&in_range) + || lp.blocks.contains(&overflow) + || in_range.uses(ctx).len() != 1 + { + return None; + } + + if const_i128(ctx, amount) != Some(1) { + return None; + } + let (lo, hi) = guard_less_than(ctx, header_term, index)?; + if counter != if subtract { hi } else { lo } { + return None; + } + Some(CheckedStep { + arm, + branch, + in_range, + in_range_operands: successor_operands(ctx, branch, 0)?, + counter, + amount, + subtract, + sum, + slot, + chain, + }) +} + +/// The pair `(lo, hi)` the header's exit test has ordered, `lo < hi`, on the +/// path into its successor `index`. `None` when the test does not order two +/// values strictly. +fn guard_less_than( + ctx: &Context, + header_term: Ptr, + index: usize, +) -> Option<(Value, Value)> { + let (test, negated) = unwrap_not(ctx, header_term.deref(ctx).get_operand(0)); + let (pred, lhs, rhs) = match_cmp(ctx, test.defining_op()?)?; + // The header enters its first successor when the condition holds. + let pred = if (index == 0) != negated { + pred + } else { + pred.negate() + }; + match pred { + CmpPred::Lt => Some((lhs, rhs)), + CmpPred::Gt => Some((rhs, lhs)), + CmpPred::Le | CmpPred::Ge => None, + } +} + +/// Replace a signed counter's checked step by a plain `add` (`sub`) and its +/// dead overflow edge by a `goto`, so the arm has the shape an unsigned counter +/// has: `arm: i1 = i + 1; goto in_range`. The tuple's operations go with the +/// edge, unless something else reads the slot. +fn fold_checked_step(ctx: &mut Context, step: CheckedStep) { + let info = if step.subtract { + MirSubOp::get_concrete_op_info() + } else { + MirAddOp::get_concrete_op_info() + }; + let ty = step.sum.get_type(ctx); + let plain = Operation::new( + ctx, + info, + vec![ty], + vec![step.counter, step.amount], + vec![], + 0, + ); + plain.insert_before(ctx, step.branch); + let next = plain.deref(ctx).get_result(0); + step.sum.replace_all_uses_with(ctx, &next); + + Operation::erase(step.branch, ctx); + let goto = Operation::new( + ctx, + MirGotoOp::get_concrete_op_info(), + vec![], + step.in_range_operands, + vec![step.in_range], + 0, + ); + goto.insert_at_back(step.arm, ctx); + + for op in step.chain.into_iter().rev() { + let dead = if Operation::get_op::(op, ctx).is_some() { + step.slot.uses(ctx).len() == 1 + } else { + op.deref(ctx) + .results() + .all(|result| result.uses(ctx).is_empty()) + }; + if dead { + Operation::erase(op, ctx); + } + } +} + /// Follow a discriminant match from `join` for a value whose discriminant is /// `discriminant`. Returns the first block the match does not decide, with the /// branch op and successor index that enter it. Every block the match passes diff --git a/cuda-oxide/crates/mir-transforms/src/forward_compiler_result_bundles.rs b/cuda-oxide/crates/mir-transforms/src/forward_compiler_result_bundles.rs index d68254639c..13ddb33274 100644 --- a/cuda-oxide/crates/mir-transforms/src/forward_compiler_result_bundles.rs +++ b/cuda-oxide/crates/mir-transforms/src/forward_compiler_result_bundles.rs @@ -48,7 +48,6 @@ use pliron::{ rewriter::{IRRewriter, Rewriter}, }, linked_list::{ContainsLinkedList, LinkedList}, - location::Located, op::Op, operation::Operation, pass::AnalysisManager, 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 cc51a82c9c..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. @@ -94,7 +96,6 @@ use pliron::{ rewriter::{IRRewriter, Rewriter}, }, linked_list::{ContainsLinkedList, LinkedList}, - location::Located, op::Op, operation::Operation, r#type::{TypeHandle, Typed}, @@ -140,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; }; @@ -151,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(); @@ -499,12 +668,10 @@ fn analyze_borrowed_pointer_read( } let function = entry_block.deref(ctx).get_parent_op(ctx)?; - Operation::get_op::(function, ctx)?; - let alwaysinline_key: pliron::identifier::Identifier = "alwaysinline".try_into().ok()?; - function - .deref(ctx) - .attributes - .get::(&alwaysinline_key)?; + let mir_function = Operation::get_op::(function, ctx)?; + if mir_function.inline_intent(ctx) != Some(dialect_mir::attributes::InlineIntentAttr::Always) { + return None; + } let aggregate_pointer_type = aggregate_pointer.get_type(ctx); let aggregate_pointer_type_ref = aggregate_pointer_type.deref(ctx); @@ -987,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(); @@ -1286,10 +1530,7 @@ mod tests { let function_op = MirFuncOp::new(ctx, function, TypeAttr::new(function_type.into())); function_op.set_symbol_name(ctx, helper_symbol.try_into().unwrap()); if alwaysinline { - function.deref_mut(ctx).attributes.set( - "alwaysinline".try_into().unwrap(), - pliron::builtin::attributes::StringAttr::new("true".to_string()), - ); + function_op.set_inline_intent(ctx, dialect_mir::attributes::InlineIntentAttr::Always); } module.append_operation(ctx, function, 0); diff --git a/cuda-oxide/crates/mir-transforms/src/unroll.rs b/cuda-oxide/crates/mir-transforms/src/unroll.rs index d3e903d21c..1616439232 100644 --- a/cuda-oxide/crates/mir-transforms/src/unroll.rs +++ b/cuda-oxide/crates/mir-transforms/src/unroll.rs @@ -20,8 +20,9 @@ //! //! The current analysis recognizes explicit counted `while` loops, including an //! exit test that adds a constant to the counter (`while i + 2 <= n`), and range -//! `for` loops (`for i in a..b`). Other iterator loops, such as `a..=b` or -//! `.step_by(k)`, are not yet recognized. +//! `for` loops (`for i in a..b` and `(a..b).rev()`) over any integer counter, +//! including `isize`. Other iterator loops, such as `a..=b` or `.step_by(k)`, +//! are not yet recognized. //! //! Several `continue` paths are supported: the pass joins their back-edges //! before unrolling. Full `#[unroll]` also preserves early `break` paths and @@ -62,6 +63,7 @@ use pliron::irbuild::{ rewriter::IRRewriter, }; use pliron::linked_list::ContainsLinkedList; +use pliron::location::{Location, Source}; use pliron::op::{Op, op_cast}; use pliron::operation::Operation; use pliron::opts::constants::sccp::sccp; @@ -111,6 +113,54 @@ fn unroll_kind(factor: u32) -> String { } } +/// A hint op's source position as `file:line:column`, the form editors and +/// terminals turn into a link. Returns `None` when the op does not carry a +/// single position in a file. +fn hint_location(ctx: &Context, op: Ptr) -> Option { + match op.deref(ctx).loc() { + Location::SrcPos { + src: Source::File(path), + pos, + } => Some(format!( + "{}:{}:{}", + pliron::uniqued_any::get(ctx, *path).display(), + pos.line, + pos.column + )), + _ => None, + } +} + +/// Unlink every hint op annotating `loop_id` and return the first located +/// one's position, for the warning that speaks for the loop. Returns `None` +/// when none of them carries a source position. +fn take_loop_hints( + ctx: &Context, + hints: &[(Ptr, Ptr, u32)], + info: &LoopInfo, + loop_id: usize, +) -> Option { + let mut location = None; + for (op, block, _f) in hints { + if info.innermost_loop(*block) == Some(loop_id) { + location = location.or_else(|| hint_location(ctx, *op)); + op.unlink(ctx); + } + } + location +} + +/// The full text of one "not unrolled" warning: the reason line, then the +/// annotation's source position on a `-->` line when `location` is `Some`. +fn not_unrolled_warning(kind: &str, reason: &str, location: Option<&str>) -> String { + let mut warning = format!("warning: {kind} requested but the loop was not unrolled: {reason}"); + if let Some(location) = location { + warning.push_str("\n --> "); + warning.push_str(location); + } + warning +} + /// Check code-growth budgets whose inputs have already been counted safely. fn check_growth_budget( copies: u64, @@ -256,8 +306,12 @@ pub fn unroll_annotated_loops( // for unrolling, so say so loudly, then drop the markers and stop. for (op, _block, factor) in &hints { eprintln!( - "warning: {} requested but the loop was not unrolled: the annotation is not inside a recognizable loop", - unroll_kind(*factor) + "{}", + not_unrolled_warning( + &unroll_kind(*factor), + "the annotation is not inside a recognizable loop", + hint_location(ctx, *op).as_deref(), + ) ); op.unlink(ctx); } @@ -267,13 +321,14 @@ pub fn unroll_annotated_loops( let Some(ph) = info.preheader(ctx, region, loop_id) else { // The author asked for unrolling; never silently do nothing. - for (op, block, _f) in &hints { - if info.innermost_loop(*block) == Some(loop_id) { - op.unlink(ctx); - } - } + let location = take_loop_hints(ctx, &hints, &info, loop_id); eprintln!( - "warning: {kind} requested but the loop was not unrolled: it has no single preheader (it is entered from more than one place)" + "{}", + not_unrolled_warning( + &kind, + "it has no single preheader (it is entered from more than one place)", + location.as_deref(), + ) ); continue; }; @@ -288,12 +343,11 @@ pub fn unroll_annotated_loops( continue; } CanonicalizeOutcome::Unsupported(reason) => { - for (op, block, _f) in &hints { - if info.innermost_loop(*block) == Some(loop_id) { - op.unlink(ctx); - } - } - eprintln!("warning: {kind} requested but the loop was not unrolled: {reason}"); + let location = take_loop_hints(ctx, &hints, &info, loop_id); + eprintln!( + "{}", + not_unrolled_warning(&kind, &reason, location.as_deref()) + ); continue; } } @@ -306,13 +360,14 @@ pub fn unroll_annotated_loops( let exiting = info.exiting_blocks(ctx, region, loop_id); let exits = info.exit_blocks(ctx, region, loop_id); if exiting.len() != 1 || exiting[0] != lp.header || exits.len() != 1 { - for (op, block, _f) in &hints { - if info.innermost_loop(*block) == Some(loop_id) { - op.unlink(ctx); - } - } + let location = take_loop_hints(ctx, &hints, &info, loop_id); eprintln!( - "warning: {kind} requested but the loop was not unrolled: partial unrolling does not yet support an early `break` or multiple exits" + "{}", + not_unrolled_warning( + &kind, + "partial unrolling does not yet support an early `break` or multiple exits", + location.as_deref(), + ) ); continue; } @@ -331,13 +386,14 @@ pub fn unroll_annotated_loops( continue; } CanonicalizeOutcome::Unsupported(reason) => { - for (op, block, _f) in &hints { - if info.innermost_loop(*block) == Some(loop_id) { - op.unlink(ctx); - } - } + let location = take_loop_hints(ctx, &hints, &info, loop_id); eprintln!( - "warning: {kind} requested but the loop was not unrolled: could not route its live-out values: {reason}" + "{}", + not_unrolled_warning( + &kind, + &format!("could not route its live-out values: {reason}"), + location.as_deref(), + ) ); continue; } @@ -356,13 +412,14 @@ pub fn unroll_annotated_loops( continue; } CanonicalizeOutcome::Unsupported(reason) => { - for (op, block, _f) in &hints { - if info.innermost_loop(*block) == Some(loop_id) { - op.unlink(ctx); - } - } + let location = take_loop_hints(ctx, &hints, &info, loop_id); eprintln!( - "warning: {kind} requested but the loop was not unrolled: could not normalize its back-edges: {reason}" + "{}", + not_unrolled_warning( + &kind, + &format!("could not normalize its back-edges: {reason}"), + location.as_deref(), + ) ); continue; } @@ -370,11 +427,7 @@ pub fn unroll_annotated_loops( // Consume this loop's hint(s) before any cloning so the markers are // never copied into the unrolled bodies. - for (op, block, _f) in &hints { - if info.innermost_loop(*block) == Some(loop_id) { - op.unlink(ctx); - } - } + let location = take_loop_hints(ctx, &hints, &info, loop_id); let rec = induction::analyze(ctx, &info, loop_id, ph); if verbose() { @@ -398,7 +451,10 @@ pub fn unroll_annotated_loops( // Requested but unsupported shape: report exactly why, loudly, so // it is never a silent no-op. UnrollOutcome::Skipped(reason) => { - eprintln!("warning: {kind} requested but the loop was not unrolled: {reason}"); + eprintln!( + "{}", + not_unrolled_warning(&kind, &reason, location.as_deref()) + ); } } } @@ -1631,7 +1687,15 @@ fn collect_hints( #[cfg(test)] mod tests { - use super::check_growth_budget; + use std::path::PathBuf; + + use super::{ + MirUnrollHintOp, check_growth_budget, hint_location, not_unrolled_warning, unroll_kind, + }; + use pliron::combine::stream::position::SourcePosition; + use pliron::context::Context; + use pliron::location::{Location, Source}; + use pliron::op::Op; #[test] fn growth_budget_counts_operations_not_only_blocks() { @@ -1642,4 +1706,72 @@ mod tests { "the exact operation limit remains allowed" ); } + + /// A `SrcPos` location like the importer attaches to a hint op. + fn src_pos(ctx: &mut Context, file: &str, line: i32, column: i32) -> Location { + Location::SrcPos { + src: Source::new_from_file(ctx, PathBuf::from(file)), + pos: SourcePosition { line, column }, + } + } + + #[test] + fn a_located_hint_op_adds_its_position_to_the_not_unrolled_warning() { + let cases = [ + ( + 0u32, + "kernel.rs", + 12, + 5, + "the annotation is not inside a recognizable loop", + ), + ( + 2u32, + "device/gemm.rs", + 3, + 27, + "partial unrolling does not yet support an early `break` or multiple exits", + ), + ]; + for (factor, file, line, column, reason) in cases { + let mut ctx = Context::new(); + let loc = src_pos(&mut ctx, file, line, column); + let hint = MirUnrollHintOp::new(&mut ctx, factor).get_operation(); + hint.deref_mut(&ctx).set_loc(loc); + + let location = hint_location(&ctx, hint).expect("a source position renders"); + let warning = not_unrolled_warning(&unroll_kind(factor), reason, Some(&location)); + let lines: Vec<&str> = warning.lines().collect(); + let case = format!("factor={factor}: {reason}"); + assert_eq!(lines.len(), 2, "{case}"); + assert!(lines[0].starts_with("warning: "), "{case}"); + assert!(lines[0].contains(reason), "{case}"); + let location_line = lines[1]; + assert!(location_line.starts_with(" --> "), "{case}"); + assert!( + location_line.ends_with(&format!("{file}:{line}:{column}")), + "{case}: {location_line}" + ); + } + } + + #[test] + fn a_hint_op_without_a_position_prints_only_the_reason_line() { + let mut ctx = Context::new(); + let hint = MirUnrollHintOp::new(&mut ctx, 0).get_operation(); + assert_eq!( + hint_location(&ctx, hint), + None, + "a fresh op has no position" + ); + + let reason = "it has no single preheader (it is entered from more than one place)"; + let warning = not_unrolled_warning( + &unroll_kind(0), + reason, + hint_location(&ctx, hint).as_deref(), + ); + assert!(!warning.contains('\n'), "{warning}"); + assert!(warning.ends_with(reason), "{warning}"); + } } diff --git a/cuda-oxide/crates/mir-transforms/tests/common/mod.rs b/cuda-oxide/crates/mir-transforms/tests/common/mod.rs index a8313be215..1fa11b9138 100644 --- a/cuda-oxide/crates/mir-transforms/tests/common/mod.rs +++ b/cuda-oxide/crates/mir-transforms/tests/common/mod.rs @@ -17,11 +17,12 @@ use core::num::NonZero; use dialect_mir::attributes::{FieldIndexAttr, MirCastKindAttr, VariantIndexAttr}; use dialect_mir::ops::{ - MirAddOp, MirCastOp, MirCondBranchOp, MirConstantOp, MirConstructEnumOp, MirEnumPayloadOp, - MirEqOp, MirFuncOp, MirGeOp, MirGetDiscriminantOp, MirGotoOp, MirGtOp, MirLeOp, MirLtOp, - MirNotOp, MirReturnOp, MirSubOp, MirUnreachableOp, + MirAddOp, MirAllocaOp, MirCastOp, MirCheckedAddOp, MirCheckedSubOp, MirCondBranchOp, + MirConstantOp, MirConstructEnumOp, MirConstructTupleOp, MirEnumPayloadOp, MirEqOp, + MirFieldAddrOp, MirFuncOp, MirGeOp, MirGetDiscriminantOp, MirGotoOp, MirGtOp, MirLeOp, + MirLoadOp, MirLtOp, MirNotOp, MirReturnOp, MirStoreOp, MirSubOp, MirUnreachableOp, }; -use dialect_mir::types::{EnumVariant, MirEnumType}; +use dialect_mir::types::{EnumVariant, MirEnumType, MirPtrType, MirTupleType}; use mir_transforms::analyses::induction::CmpPred; use pliron::basic_block::BasicBlock; use pliron::builtin::attributes::{IntegerAttr, TypeAttr}; @@ -1110,6 +1111,24 @@ pub enum RangeBound { Param, } +/// How the `Some` arm of a [`range_for_loop`] computes the next counter. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StepForm { + /// `mir.add` (`mir.sub` for `rev`): what an unsigned counter's + /// `unchecked_add` imports as. + Plain, + /// What a signed counter's `forward_unchecked` leaves after mem2reg: + /// `mir.checked_add` (`mir.checked_sub`) whose tuple goes through a stack + /// slot that mem2reg cannot promote, and a branch on the overflow flag to a + /// block that ends in `mir.unreachable`. + Checked, + /// [`StepForm::Checked`], but the overflow edge returns to the loop header. + CheckedOverflowReenters, + /// [`StepForm::Checked`] stepping by 2, for which `i < end` does not rule + /// out overflow. + CheckedStepTwo, +} + /// The variations of [`range_for_loop`]. #[derive(Clone, Copy, Debug)] pub struct RangeForShape { @@ -1122,6 +1141,8 @@ pub struct RangeForShape { /// The exit returns `acc + next`, a value from the match that both the /// break and the normal exit reach. pub exit_reads_next: bool, + /// How the counter steps. + pub step_form: StepForm, } impl RangeForShape { @@ -1133,6 +1154,7 @@ impl RangeForShape { rev: false, break_at: None, exit_reads_next: false, + step_form: StepForm::Plain, } } } @@ -1186,6 +1208,18 @@ fn construct_enum( /// is live across the match. With `rev` the counter starts at `end`, the header tests `start < i`, and /// `some` yields and passes on `i - 1`. `break_at` adds /// `cond_br x < k [latch, exit(acc1)]` after the addition. +/// +/// With a [`StepForm`] other than `Plain`, `some` is the shape of a signed +/// counter, as the backend dumps it: +/// +/// ```text +/// entry: slot = alloca (ty, i1) +/// some: t = checked_add(i, 1); store slot, t +/// x = load field_addr(slot, 0); f = load field_addr(slot, 1) +/// cond_br not(f) [ok, overflow] +/// ok: o = Some(i); goto join(o, x) +/// overflow: construct_tuple (); goto unreachable +/// ``` pub fn range_for_loop( ctx: &mut Context, ty: TypedHandle, @@ -1234,6 +1268,20 @@ pub fn range_for_loop( let end = bound(ctx, shape.end); let acc0 = iconst(ctx, entry, ty, 0); let i0 = if shape.rev { end } else { start }; + let pair: TypeHandle = MirTupleType::get(ctx, vec![ty.into(), i1.into()]).into(); + let slot = (shape.step_form != StepForm::Plain).then(|| { + let slot_ptr: TypeHandle = MirPtrType::get_generic(ctx, pair, true).into(); + let alloca = Operation::new( + ctx, + MirAllocaOp::get_concrete_op_info(), + vec![slot_ptr], + vec![], + vec![], + 0, + ); + alloca.insert_at_back(entry, ctx); + alloca.deref(ctx).get_result(0) + }); goto(ctx, entry, header, vec![acc0, i0]); let acc = header.deref(ctx).get_argument(0); @@ -1261,16 +1309,92 @@ pub fn range_for_loop( }; cond_br(ctx, header, done, none, some); - let one = iconst(ctx, some, ty, 1); - let step = if shape.rev { - MirSubOp::get_concrete_op_info() + let amount = if shape.step_form == StepForm::CheckedStepTwo { + 2 + } else { + 1 + }; + let one = iconst(ctx, some, ty, amount); + // The block that builds `Some` and goes to `join`. + let mut arm = some; + let inext = if let Some(slot) = slot { + let checked = if shape.rev { + MirCheckedSubOp::get_concrete_op_info() + } else { + MirCheckedAddOp::get_concrete_op_info() + }; + let pair_value = op2!(ctx, some, checked, pair, i, one); + let store = Operation::new( + ctx, + MirStoreOp::get_concrete_op_info(), + vec![], + vec![slot, pair_value], + vec![], + 0, + ); + store.insert_at_back(some, ctx); + let load_field = |ctx: &mut Context, field: u32, field_ty: TypeHandle| { + let field_ptr: TypeHandle = MirPtrType::get_generic(ctx, field_ty, true).into(); + let address = MirFieldAddrOp::build(ctx, slot, field_ptr, field) + .expect("the slot is a MIR pointer"); + address.insert_at_back(some, ctx); + let field_value = address.deref(ctx).get_result(0); + let load = Operation::new( + ctx, + MirLoadOp::get_concrete_op_info(), + vec![field_ty], + vec![field_value], + vec![], + 0, + ); + load.insert_at_back(some, ctx); + load.deref(ctx).get_result(0) + }; + let sum = load_field(ctx, 0, ty.into()); + let overflowed = load_field(ctx, 1, i1.into()); + let in_range = { + let op = Operation::new( + ctx, + MirNotOp::get_concrete_op_info(), + vec![i1.into()], + vec![overflowed], + vec![], + 0, + ); + op.insert_at_back(some, ctx); + op.deref(ctx).get_result(0) + }; + let ok = block(ctx, region, vec![]); + let overflow = block(ctx, region, vec![]); + cond_br(ctx, some, in_range, ok, overflow); + let unit: TypeHandle = MirTupleType::get(ctx, vec![]).into(); + Operation::new( + ctx, + MirConstructTupleOp::get_concrete_op_info(), + vec![unit], + vec![], + vec![], + 0, + ) + .insert_at_back(overflow, ctx); + if shape.step_form == StepForm::CheckedOverflowReenters { + goto(ctx, overflow, header, vec![acc, i]); + } else { + goto(ctx, overflow, unreachable, vec![]); + } + arm = ok; + sum } else { - MirAddOp::get_concrete_op_info() + let step = if shape.rev { + MirSubOp::get_concrete_op_info() + } else { + MirAddOp::get_concrete_op_info() + }; + op2!(ctx, some, step, ty.into(), i, one) }; - let inext = op2!(ctx, some, step, ty.into(), i, one); let yielded = if shape.rev { inext } else { i }; - let some_value = construct_enum(ctx, some, option, 1, vec![yielded]); - goto(ctx, some, join, vec![some_value, inext]); + let some_value = construct_enum(ctx, arm, option, 1, vec![yielded]); + goto(ctx, arm, join, vec![some_value, inext]); let none_value = construct_enum(ctx, none, option, 0, vec![]); goto(ctx, none, join, vec![none_value, i]); diff --git a/cuda-oxide/crates/mir-transforms/tests/unroll.rs b/cuda-oxide/crates/mir-transforms/tests/unroll.rs index 00534a4056..032c928c32 100644 --- a/cuda-oxide/crates/mir-transforms/tests/unroll.rs +++ b/cuda-oxide/crates/mir-transforms/tests/unroll.rs @@ -14,14 +14,15 @@ mod common; use common::{ - OffsetBound, RangeBound, RangeForLoop, RangeForShape, counted_loop, counted_loop_from_step, - early_exit_counted_loop, early_exit_with_direct_liveout, mir_ctx, multi_latch_counted_loop, - multiple_exit_counted_loop, nested_counted_loop, offset_counted_loop, range_for_loop, u32t, + OffsetBound, RangeBound, RangeForLoop, RangeForShape, StepForm, counted_loop, + counted_loop_from_step, early_exit_counted_loop, early_exit_with_direct_liveout, i32t, mir_ctx, + multi_latch_counted_loop, multiple_exit_counted_loop, nested_counted_loop, offset_counted_loop, + range_for_loop, u32t, }; use dialect_mir::ops::{ - MirAddOp, MirBitAndOp, MirCallOp, MirCondBranchOp, MirConstantOp, MirGeOp, - MirGetDiscriminantOp, MirGtOp, MirLeOp, MirLtOp, MirNotOp, MirReturnOp, MirSubOp, - MirUnrollHintOp, + MirAddOp, MirBitAndOp, MirCallOp, MirCheckedAddOp, MirCheckedSubOp, MirCondBranchOp, + MirConstantOp, MirGeOp, MirGetDiscriminantOp, MirGtOp, MirLeOp, MirLtOp, MirNotOp, MirReturnOp, + MirSubOp, MirUnrollHintOp, }; use mir_transforms::analyses::induction::{CmpPred, analyze}; use mir_transforms::unroll::unroll_annotated_loops; @@ -995,6 +996,138 @@ fn partial_unroll_of_a_range_for_loop_returns_the_wrapping_sum_for_every_end() { } } +/// The 8-bit signed bounds a range loop over a checked-step counter is tested at. +fn signed_eight_bit_case() -> (i128, i128, Vec) { + let (_, min, max, bounds) = eight_bit_cases() + .into_iter() + .find(|case| case.0 == Signedness::Signed) + .expect("a signed case"); + (min, max, bounds) +} + +#[test] +fn full_unroll_of_a_signed_range_for_loop_folds_the_wrapping_sum_for_every_boundary_range() { + let (_, _, bounds) = signed_eight_bit_case(); + for rev in [false, true] { + for &start in &bounds { + for &end in &bounds { + let mut ctx = mir_ctx(); + let ty = IntegerType::get(&ctx, 8, Signedness::Signed); + let shape = RangeForShape { + rev, + step_form: StepForm::Checked, + ..RangeForShape::new( + RangeBound::Const(start as i64), + RangeBound::Const(end as i64), + ) + }; + let lp = range_for_loop(&mut ctx, ty, shape); + unroll_range_for(&mut ctx, &lp, 0); + let case = format!("{start}..{end}, rev={rev}"); + assert_eq!(loop_count(&ctx, lp.region), 0, "{case}"); + assert_eq!( + sole_return_constant(&ctx, lp.region).map(|v| v.rem_euclid(256)), + Some(wrapping_range_sum(Signedness::Signed, start, end).rem_euclid(256)), + "{case}" + ); + assert_eq!( + op_count::(&ctx, lp.region) + + op_count::(&ctx, lp.region), + 0, + "{case}: the overflow check is gone with the dead edge" + ); + } + } + } +} + +#[test] +fn partial_unroll_of_a_signed_range_for_loop_returns_the_wrapping_sum_for_every_end() { + let (min, max, bounds) = signed_eight_bit_case(); + for factor in [2, 3, 4] { + for &start in &bounds { + let mut ctx = mir_ctx(); + let ty = IntegerType::get(&ctx, 8, Signedness::Signed); + let shape = RangeForShape { + step_form: StepForm::Checked, + ..RangeForShape::new(RangeBound::Const(start as i64), RangeBound::Param) + }; + let lp = range_for_loop(&mut ctx, ty, shape); + unroll_range_for(&mut ctx, &lp, factor); + let case = format!("{start}..n, factor={factor}"); + assert_eq!(loop_count(&ctx, lp.region), 2, "{case}"); + for end in min..=max { + assert_eq!( + run_function(&ctx, lp.region, &[end]), + wrapping_range_sum(Signedness::Signed, start, end), + "{case}, n={end}" + ); + } + } + } +} + +#[test] +fn signed_range_for_loops_the_unroller_refuses_keep_their_loop_and_results() { + // The checked step is folded first; the unroller then refuses these loops. + let runtime_end = RangeForShape::new(RangeBound::Const(0), RangeBound::Param); + let runtime_start = RangeForShape::new(RangeBound::Param, RangeBound::Const(8)); + for (shape, factor) in [(runtime_end, 0), (runtime_start, 0)] { + for rev in [false, true] { + let shape = RangeForShape { + rev, + step_form: StepForm::Checked, + ..shape + }; + let mut ctx = mir_ctx(); + let i32 = i32t(&mut ctx); + let lp = range_for_loop(&mut ctx, i32, shape); + unroll_range_for(&mut ctx, &lp, factor); + assert_eq!(loop_count(&ctx, lp.region), 1, "{shape:?}"); + for arg in [0, 1, 4, 9] { + let (start, end) = match shape.start { + RangeBound::Param => (arg, 8), + RangeBound::Const(_) => (0, arg), + }; + assert_eq!( + run_function(&ctx, lp.region, &[arg]), + (start..end).sum::(), + "{shape:?}, argument {arg}" + ); + } + } + } +} + +#[test] +fn checked_step_whose_overflow_edge_is_not_dead_keeps_the_loop_untouched() { + for step_form in [StepForm::CheckedOverflowReenters, StepForm::CheckedStepTwo] { + for rev in [false, true] { + let mut ctx = mir_ctx(); + let i32 = i32t(&mut ctx); + let shape = RangeForShape { + rev, + step_form, + ..RangeForShape::new(RangeBound::Const(0), RangeBound::Const(8)) + }; + let lp = range_for_loop(&mut ctx, i32, shape); + let before = operations(&ctx, lp.region).len(); + + unroll_range_for(&mut ctx, &lp, 0); + + let case = format!("{step_form:?}, rev={rev}"); + assert_eq!(loop_count(&ctx, lp.region), 1, "{case}"); + assert_eq!(operations(&ctx, lp.region).len(), before, "{case}"); + assert_eq!( + op_count::(&ctx, lp.region) + + op_count::(&ctx, lp.region), + 1, + "{case}: the overflow check is still there" + ); + } + } +} + #[test] fn full_unroll_keeps_a_break_in_a_range_for_loop() { let mut ctx = mir_ctx(); diff --git a/cuda-oxide/crates/nvjitlink-sys/README.md b/cuda-oxide/crates/nvjitlink-sys/README.md index 8368fd4e13..446b16a1a2 100644 --- a/cuda-oxide/crates/nvjitlink-sys/README.md +++ b/cuda-oxide/crates/nvjitlink-sys/README.md @@ -18,10 +18,10 @@ None. The library is loaded at runtime, so the CUDA Toolkit only needs to be pre `LibNvJitLink::load()` tries (in order): 1. `LIBNVJITLINK_PATH` env var, if set. -2. `/lib64/libnvJitLink.so` for `` in `CUDA_TOOLKIT_PATH`, `CUDA_HOME`, `CUDA_PATH`, `/usr/local/cuda`, `/opt/cuda`. +2. `/lib64/libnvJitLink.so`, then `/lib/libnvJitLink.so`, for `` in `CUDA_TOOLKIT_PATH`, `CUDA_HOME`, `CUDA_PATH`, `/usr/local/cuda`, `/opt/cuda`. 3. The system loader (`libnvJitLink.so.13`, `libnvJitLink.so.12`, `libnvJitLink.so`). -nvJitLink ships with the standard CUDA Toolkit at `/lib64/`. No separate download. +nvJitLink ships with the standard CUDA Toolkit at `/lib64/`; conda-forge and nixpkgs layouts place it at `/lib/`. No separate download. Only the unversioned `libnvJitLink.so` name is probed under a root. A layout that ships only `libnvJitLink.so.13` is found through the system loader instead, and such a load cannot be tied to one file on disk, which cubin materialization (`cargo oxide build --materialize-cubin`) requires. Set `LIBNVJITLINK_PATH` to the versioned file in that case. ## Symbol naming diff --git a/cuda-oxide/crates/nvjitlink-sys/src/lib.rs b/cuda-oxide/crates/nvjitlink-sys/src/lib.rs index 4ff5f279fc..c70606565d 100644 --- a/cuda-oxide/crates/nvjitlink-sys/src/lib.rs +++ b/cuda-oxide/crates/nvjitlink-sys/src/lib.rs @@ -7,7 +7,24 @@ //! //! nvJitLink links one or more LTOIR modules (and other input forms) into //! a final cubin or PTX. It is part of the CUDA Toolkit and ships at -//! `/lib64/libnvJitLink.so`. +//! `/lib64/libnvJitLink.so`. conda-forge and nixpkgs lay the toolkit +//! out with `lib/` instead of `lib64/`, so both directories are probed. +//! +//! # Library discovery +//! +//! [`LibNvJitLink::load`] tries (in order): +//! 1. `LIBNVJITLINK_PATH` env var, if set. +//! 2. `/lib64/libnvJitLink.so`, then `/lib/libnvJitLink.so`, for +//! `` in `CUDA_TOOLKIT_PATH`, `CUDA_HOME`, `CUDA_PATH`, +//! `/usr/local/cuda`, `/opt/cuda`. +//! 3. The system loader (`libnvJitLink.so.13`, `libnvJitLink.so.12`, +//! `libnvJitLink.so`). +//! +//! Only the unversioned `libnvJitLink.so` name is probed under a root. A +//! layout that ships only `libnvJitLink.so.13` is reached through the system +//! loader, which does not say which file it mapped, so that load cannot be +//! tied to one file on disk. Set `LIBNVJITLINK_PATH` to the file when the +//! exact file matters, as it does for cubin materialization. //! //! # Symbol naming //! @@ -729,12 +746,8 @@ fn open_library(tried: &mut Vec, retain_exact_file: bool) -> Option, retain_exact_file: bool) -> Option, + retain_exact_file: bool, +) -> Option { + for path in library_candidates(roots) { + tried.push(path.display().to_string()); + if let Some(opened) = open_library_path(&path, retain_exact_file) { + return Some(opened); + } + } + None +} + fn open_library_path(path: &Path, retain_exact_file: bool) -> Option { #[cfg(not(target_os = "linux"))] let _ = retain_exact_file; @@ -792,6 +824,21 @@ fn open_library_path(path: &Path, retain_exact_file: bool) -> Option Vec { + roots + .iter() + .flat_map(|root| ["lib64", "lib"].map(|dir| root.join(dir).join("libnvJitLink.so"))) + .collect() +} + fn cuda_roots() -> Vec { cuda_roots_from_env(|var| std::env::var(var).ok()) } @@ -922,6 +969,102 @@ mod tests { std::fs::remove_dir_all(directory).unwrap(); } + /// A fresh directory under the system temp dir, used as a CUDA root. + #[cfg(target_os = "linux")] + fn scratch_root(tag: &str) -> PathBuf { + let nonce = SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let root = std::env::temp_dir().join(format!( + "nvjitlink-sys-{tag}-{}-{nonce}", + std::process::id() + )); + std::fs::create_dir(&root).unwrap(); + root + } + + /// Build the probe DSO as `//libnvJitLink.so`. + #[cfg(target_os = "linux")] + fn place_probe_library(root: &Path, dir: &str, value: i32) -> PathBuf { + let directory = root.join(dir); + std::fs::create_dir(&directory).unwrap(); + let library_path = directory.join("libnvJitLink.so"); + compile_probe_library(&directory.join("probe.c"), &library_path, value); + library_path + } + + #[cfg(target_os = "linux")] + fn probe_value(opened: &OpenedLibrary) -> c_int { + let probe: Symbol c_int> = + unsafe { opened.library.get(b"cuda_oxide_probe") }.unwrap(); + unsafe { probe() } + } + + #[cfg(target_os = "linux")] + fn file_identity(path: &Path) -> Option { + LibraryFileIdentity::capture_file(&File::open(path).unwrap()) + } + + #[cfg(target_os = "linux")] + fn candidate_strings(root: &Path, dirs: &[&str]) -> Vec { + dirs.iter() + .map(|dir| root.join(dir).join("libnvJitLink.so").display().to_string()) + .collect() + } + + #[cfg(target_os = "linux")] + #[test] + fn root_probe_loads_lib_when_lib64_is_absent() { + let root = scratch_root("root-lib"); + let library_path = place_probe_library(&root, "lib", 7); + let expected_identity = file_identity(&library_path); + assert!(expected_identity.is_some()); + + let mut tried = Vec::new(); + let opened = open_library_under_roots(std::slice::from_ref(&root), &mut tried, true) + .expect("lib/libnvJitLink.so is found when lib64/ is absent"); + + assert_eq!(probe_value(&opened), 7); + assert_eq!(opened.loaded_identity, expected_identity); + assert_eq!(tried, candidate_strings(&root, &["lib64", "lib"])); + + drop(opened); + std::fs::remove_dir_all(root).unwrap(); + } + + #[cfg(target_os = "linux")] + #[test] + fn root_probe_prefers_lib64_over_lib() { + let root = scratch_root("root-both"); + let lib64_path = place_probe_library(&root, "lib64", 1); + place_probe_library(&root, "lib", 2); + let lib64_identity = file_identity(&lib64_path); + + let mut tried = Vec::new(); + let opened = open_library_under_roots(std::slice::from_ref(&root), &mut tried, true) + .expect("lib64/libnvJitLink.so is found"); + + assert_eq!(probe_value(&opened), 1); + assert_eq!(opened.loaded_identity, lib64_identity); + assert_eq!(tried, candidate_strings(&root, &["lib64"])); + + drop(opened); + std::fs::remove_dir_all(root).unwrap(); + } + + #[cfg(target_os = "linux")] + #[test] + fn root_probe_reports_both_candidates_for_an_empty_root() { + let root = scratch_root("root-empty"); + + let mut tried = Vec::new(); + assert!(open_library_under_roots(std::slice::from_ref(&root), &mut tried, true).is_none()); + assert_eq!(tried, candidate_strings(&root, &["lib64", "lib"])); + + std::fs::remove_dir_all(root).unwrap(); + } + #[test] fn open_descriptor_remains_bound_to_replaced_inode() { let nonce = SystemTime::now() @@ -1082,6 +1225,20 @@ mod tests { assert_eq!(prepared.as_ref(), auto_detected); } + #[test] + fn library_candidates_try_lib64_then_lib_under_each_root() { + let roots = [PathBuf::from("/cuda"), PathBuf::from("/conda")]; + assert_eq!( + library_candidates(&roots), + vec![ + PathBuf::from("/cuda/lib64/libnvJitLink.so"), + PathBuf::from("/cuda/lib/libnvJitLink.so"), + PathBuf::from("/conda/lib64/libnvJitLink.so"), + PathBuf::from("/conda/lib/libnvJitLink.so"), + ] + ); + } + #[test] fn cuda_roots_prefers_project_toolkit_env_var() { let roots = cuda_roots_from_env(|var| match var { diff --git a/cuda-oxide/crates/nvvm-transforms/src/legalize.rs b/cuda-oxide/crates/nvvm-transforms/src/legalize.rs index f7c1768721..a3fb9f3e79 100644 --- a/cuda-oxide/crates/nvvm-transforms/src/legalize.rs +++ b/cuda-oxide/crates/nvvm-transforms/src/legalize.rs @@ -27,7 +27,7 @@ use llvm_export::{ use pliron::{ attribute::Attribute, builtin::{ - attributes::{BoolAttr, FPDoubleAttr, FPSingleAttr, IntegerAttr, TypeAttr}, + attributes::{BoolAttr, FPDoubleAttr, FPSingleAttr, IntegerAttr, StringAttr, TypeAttr}, op_interfaces::{CallOpCallable, CallOpInterface, SymbolOpInterface}, type_interfaces::FunctionTypeInterface, types::{FP16Type, FP32Type, FP64Type, IntegerType, Signedness}, @@ -35,7 +35,6 @@ use pliron::{ context::{Context, Ptr}, identifier::Identifier, linked_list::ContainsLinkedList, - location::Located, op::{Op, op_cast}, operation::Operation, result::Result, @@ -210,6 +209,205 @@ pub(crate) fn legalize_nvvm_bit_intrinsics( Ok(()) } +/// A function definition whose return type is a bare pointer in a non-generic +/// address space, with the type it returned before legalization. +struct PointerReturnRewrite { + func: llvm::FuncOp, + name: Identifier, + original_return: TypeHandle, +} + +/// Give every non-kernel function definition that returns a bare pointer in a +/// non-generic address space a generic pointer return instead. +/// +/// libNVVM declares such a return wrongly. Under the NVVM data layout a +/// shared-memory pointer is 32 bits wide, and libNVVM prints the callee's +/// return slot as `.param .b64 func_retval0` while the body stores 32 bits +/// into it and every call site declares a 32-bit `retval0`; ptxas then +/// rejects the module. The route never met this while such helpers were +/// inlined before reaching libNVVM, but a `noinline` helper returning +/// `*mut SharedArray` is a real function. A generic pointer is 64 bits on both +/// sides, so the function returns one and each side of the call restores the +/// address space with an `addrspacecast`: the callee casts to generic right +/// before each `ret`, and each call site casts the generic result back to the +/// original address space, so every use keeps the type it was lowered with. +/// +/// Parameters, aggregates holding such pointers and kernels are unaffected and +/// left alone, as are declarations: their ABI belongs to code cuda-oxide did +/// not compile. cuda-oxide emits both sides of every call to a definition, so +/// its signature can change; a function whose address is taken is reported +/// instead, since an indirect caller would still expect the old signature. +pub(crate) fn legalize_pointer_returns(ctx: &mut Context, module: Ptr) -> Result<()> { + let rewrites = pointer_return_rewrites(ctx, module); + if rewrites.is_empty() { + return Ok(()); + } + + let mut ops = Vec::new(); + collect_ops(ctx, module, &mut ops); + + // Validate before changing anything. + for &op in &ops { + let Some(address_of) = Operation::get_op::(op, ctx) else { + continue; + }; + let name = address_of.get_global_name(ctx); + if rewrites.iter().any(|rewrite| rewrite.name == name) { + return pliron::input_err!( + op.deref(ctx).loc(), + "cannot legalize the non-generic pointer return of @{name} for libNVVM: its address is taken" + ); + } + } + + let generic_pointer: TypeHandle = + llvm_types::PointerType::get(ctx, llvm_types::address_space::GENERIC).into(); + for rewrite in &rewrites { + retype_pointer_return(ctx, rewrite, generic_pointer)?; + } + for op in ops { + let Some(call) = Operation::get_op::(op, ctx) else { + continue; + }; + let CallOpCallable::Direct(callee) = call.callee(ctx) else { + continue; + }; + let Some(rewrite) = rewrites.iter().find(|rewrite| rewrite.name == callee) else { + continue; + }; + cast_call_result_back(ctx, op, &call, rewrite)?; + } + Ok(()) +} + +fn pointer_return_rewrites(ctx: &Context, module: Ptr) -> Vec { + let kernel_key: Identifier = "gpu_kernel" + .try_into() + .expect("static attribute name is valid"); + let Some(region) = module.deref(ctx).regions().next() else { + return Vec::new(); + }; + let Some(block) = region.deref(ctx).iter(ctx).next() else { + return Vec::new(); + }; + let top_level: Vec<_> = block.deref(ctx).iter(ctx).collect(); + let mut rewrites = Vec::new(); + for op in top_level { + let Some(func) = Operation::get_op::(op, ctx) else { + continue; + }; + // A declaration names a function cuda-oxide did not compile. + if func.get_entry_block(ctx).is_none() { + continue; + } + // A kernel entry returns nothing and is never called. + if op + .deref(ctx) + .attributes + .get::(&kernel_key) + .is_some() + { + continue; + } + let original_return = func.get_type(ctx).deref(ctx).result_type(); + let Some(space) = pointer_address_space(ctx, original_return) else { + continue; + }; + if space == llvm_types::address_space::GENERIC { + continue; + } + rewrites.push(PointerReturnRewrite { + func, + name: func.get_symbol_name(ctx), + original_return, + }); + } + rewrites +} + +fn pointer_address_space(ctx: &Context, ty: TypeHandle) -> Option { + ty.deref(ctx) + .downcast_ref::() + .map(llvm_types::PointerType::address_space) +} + +/// Change the function's return type to a generic pointer and cast each +/// returned value to generic right before its `ret`. +fn retype_pointer_return( + ctx: &mut Context, + rewrite: &PointerReturnRewrite, + generic_pointer: TypeHandle, +) -> Result<()> { + let (arg_types, is_var_arg) = { + let func_type = rewrite.func.get_type(ctx); + let func_type = func_type.deref(ctx); + (func_type.arg_types(), func_type.is_var_arg()) + }; + let new_type = llvm_types::FuncType::get(ctx, generic_pointer, arg_types, is_var_arg); + rewrite + .func + .set_attr_llvm_func_type(ctx, TypeAttr::new(new_type.into())); + + let mut body = Vec::new(); + collect_ops(ctx, rewrite.func.get_operation(), &mut body); + for op in body { + let Some(ret) = Operation::get_op::(op, ctx) else { + continue; + }; + let Some(value) = ret.retval(ctx) else { + return pliron::input_err!( + op.deref(ctx).loc(), + "@{} returns a pointer, but this `ret` carries no value", + rewrite.name + ); + }; + let loc = op.deref(ctx).loc().clone(); + let cast = llvm::AddrSpaceCastOp::new(ctx, value, generic_pointer); + cast.get_operation().deref_mut(ctx).set_loc(loc); + let generic = insert_before(ctx, op, cast.get_operation()); + Operation::replace_operand(op, ctx, 0, generic); + } + Ok(()) +} + +/// Replace a direct call to a retyped function with one that returns the +/// generic pointer, cast back to the address space its uses were lowered with. +fn cast_call_result_back( + ctx: &mut Context, + old_call: Ptr, + call: &llvm::CallOp, + rewrite: &PointerReturnRewrite, +) -> Result<()> { + let args = call.args(ctx); + let callee_type = rewrite.func.get_type(ctx); + let new_call = llvm::CallOp::new( + ctx, + CallOpCallable::Direct(rewrite.name.clone()), + callee_type, + args, + ); + let new_op = new_call.get_operation(); + { + // Keep the call's own attributes (fast-math flags, call-site + // attributes) and its location. + let attributes = old_call.deref(ctx).attributes.clone(); + let loc = old_call.deref(ctx).loc().clone(); + let mut new_call_ref = new_op.deref_mut(ctx); + new_call_ref.attributes = attributes; + new_call_ref.set_loc(loc); + } + // The copied dictionary still named the old callee type. + let callee_type = rewrite.func.get_type(ctx); + new_call.set_callee_type(ctx, callee_type.into()); + let generic = insert_before(ctx, old_call, new_op); + + let loc = old_call.deref(ctx).loc().clone(); + let cast = llvm::AddrSpaceCastOp::new(ctx, generic, rewrite.original_return); + cast.get_operation().deref_mut(ctx).set_loc(loc); + let restored = insert_before(ctx, old_call, cast.get_operation()); + replace_one_result(ctx, old_call, restored) +} + /// CUDA 12's LLVM 7 NVVM dialect does not support the scalar `half` type. /// /// This checks operands, results, function and global signatures, block @@ -2016,6 +2214,231 @@ mod tests { .collect() } + /// `define ptr addrspace(3) @recover(i64)`: the address comes in as an + /// integer, becomes a generic pointer and is cast to shared before `ret`. + fn shared_pointer_recover_function( + ctx: &mut Context, + module: &ModuleOp, + ) -> (llvm::FuncOp, Ptr) { + let i64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Signless).into(); + let shared_ptr: TypeHandle = PointerType::get(ctx, 3).into(); + let generic_ptr: TypeHandle = PointerType::get(ctx, 0).into(); + let (recover, entry) = function(ctx, module, "recover", shared_ptr, vec![i64_ty]); + let address = entry.deref(ctx).get_argument(0); + let generic = llvm::IntToPtrOp::new(ctx, address, generic_ptr); + generic.get_operation().insert_at_back(entry, ctx); + let generic_value = generic.get_operation().deref(ctx).get_result(0); + let shared = llvm::AddrSpaceCastOp::new(ctx, generic_value, shared_ptr); + shared.get_operation().insert_at_back(entry, ctx); + let shared_value = shared.get_operation().deref(ctx).get_result(0); + llvm::ReturnOp::new(ctx, Some(shared_value)) + .get_operation() + .insert_at_back(entry, ctx); + (recover, entry) + } + + /// `define void @use_shared(i64, ptr)`: calls `@recover`, loads a float + /// through the shared pointer it returns and stores it to the generic one. + fn shared_pointer_caller( + ctx: &mut Context, + module: &ModuleOp, + recover: &llvm::FuncOp, + ) -> (Ptr, Ptr) { + let i64_ty: TypeHandle = IntegerType::get(ctx, 64, Signedness::Signless).into(); + let generic_ptr: TypeHandle = PointerType::get(ctx, 0).into(); + let f32_ty: TypeHandle = FP32Type::get(ctx).into(); + let void_ty: TypeHandle = VoidType::get(ctx).into(); + let (_user, entry) = function( + ctx, + module, + "use_shared", + void_ty, + vec![i64_ty, generic_ptr], + ); + let (address, out) = { + let block = entry.deref(ctx); + (block.get_argument(0), block.get_argument(1)) + }; + let callee_type = recover.get_type(ctx); + let call = llvm::CallOp::new( + ctx, + CallOpCallable::Direct("recover".try_into().unwrap()), + callee_type, + vec![address], + ); + call.get_operation().insert_at_back(entry, ctx); + let pointer = call.get_operation().deref(ctx).get_result(0); + let load = llvm::LoadOp::new(ctx, pointer, f32_ty); + load.get_operation().insert_at_back(entry, ctx); + let value = load.get_operation().deref(ctx).get_result(0); + llvm::StoreOp::new(ctx, value, out) + .get_operation() + .insert_at_back(entry, ctx); + llvm::ReturnOp::new(ctx, None) + .get_operation() + .insert_at_back(entry, ctx); + (call.get_operation(), load.get_operation()) + } + + #[test] + fn non_generic_pointer_returns_become_generic_with_casts_on_both_sides() { + for dialect in [NvvmIrDialect::Modern, NvvmIrDialect::LegacyLlvm7] { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "shared_return".try_into().unwrap()); + let (recover, _) = shared_pointer_recover_function(&mut ctx, &module); + let (_call, load) = shared_pointer_caller(&mut ctx, &module, &recover); + let shared_ptr: TypeHandle = PointerType::get(&ctx, 3).into(); + let generic_ptr: TypeHandle = PointerType::get(&ctx, 0).into(); + module.get_operation().deref(&ctx).verify(&ctx).unwrap(); + + crate::legalize_for_nvvm(&mut ctx, module.get_operation(), dialect, 90).unwrap(); + + module.get_operation().deref(&ctx).verify(&ctx).unwrap(); + assert_eq!( + recover.get_type(&ctx).deref(&ctx).result_type(), + generic_ptr, + "{dialect:?}: the definition must return a generic pointer" + ); + assert_eq!( + load.deref(&ctx).get_operand(0).get_type(&ctx), + shared_ptr, + "{dialect:?}: the caller's load must keep its shared pointer" + ); + let ids = operation_ids(&ctx, module.get_operation()); + assert_eq!( + ids.iter().filter(|id| *id == "llvm.addrspacecast").count(), + 3, + "{dialect:?}: one cast before the ret and one after the call join the original: {ids:?}" + ); + + if dialect == NvvmIrDialect::Modern { + let ir = export_module_to_string_with_config( + &ctx, + &module, + &NvvmExportConfig::new(NvvmIrDialect::Modern), + ) + .expect("legalized module exports as modern NVVM IR"); + assert!(ir.contains("define ptr @recover(i64 "), "{ir}"); + assert!(!ir.contains("ptr addrspace(3) @recover"), "{ir}"); + assert!(ir.contains("ret ptr %"), "{ir}"); + assert!(ir.contains("addrspacecast ptr addrspace(3) %"), "{ir}"); + assert!(ir.contains("call ptr @recover(i64 "), "{ir}"); + assert_eq!( + ir.matches("to ptr addrspace(3)").count(), + 2, + "the call site casts back to shared next to the original cast:\n{ir}" + ); + } + } + } + + #[test] + fn pointer_return_legalization_leaves_other_signatures_alone() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "untouched".try_into().unwrap()); + let i64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Signless).into(); + let shared_ptr: TypeHandle = PointerType::get(&ctx, 3).into(); + let generic_ptr: TypeHandle = PointerType::get(&ctx, 0).into(); + let void_ty: TypeHandle = VoidType::get(&ctx).into(); + + // A generic pointer return needs nothing. + let (generic_return, entry) = function( + &mut ctx, + &module, + "generic_return", + generic_ptr, + vec![i64_ty], + ); + let address = entry.deref(&ctx).get_argument(0); + let generic = llvm::IntToPtrOp::new(&mut ctx, address, generic_ptr); + generic.get_operation().insert_at_back(entry, &ctx); + let generic_value = generic.get_operation().deref(&ctx).get_result(0); + llvm::ReturnOp::new(&mut ctx, Some(generic_value)) + .get_operation() + .insert_at_back(entry, &ctx); + + // A shared pointer parameter is declared correctly by libNVVM. + let (shared_param, entry) = + function(&mut ctx, &module, "shared_param", void_ty, vec![shared_ptr]); + llvm::ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + + // A declaration belongs to code cuda-oxide did not compile. + let declared_type = FuncType::get(&ctx, shared_ptr, vec![i64_ty], false); + let declared = llvm::FuncOp::new(&mut ctx, "external".try_into().unwrap(), declared_type); + declared + .get_operation() + .insert_at_back(module_block(&ctx, &module), &ctx); + let (_caller, entry) = function(&mut ctx, &module, "caller", void_ty, vec![i64_ty]); + let address = entry.deref(&ctx).get_argument(0); + let call = llvm::CallOp::new( + &mut ctx, + CallOpCallable::Direct("external".try_into().unwrap()), + declared_type, + vec![address], + ); + call.get_operation().insert_at_back(entry, &ctx); + llvm::ReturnOp::new(&mut ctx, None) + .get_operation() + .insert_at_back(entry, &ctx); + + let before = operation_ids(&ctx, module.get_operation()); + crate::legalize_for_nvvm(&mut ctx, module.get_operation(), NvvmIrDialect::Modern, 120) + .unwrap(); + + assert_eq!(operation_ids(&ctx, module.get_operation()), before); + assert_eq!( + generic_return.get_type(&ctx).deref(&ctx).result_type(), + generic_ptr + ); + assert_eq!( + shared_param.get_type(&ctx).deref(&ctx).arg_types(), + vec![shared_ptr] + ); + assert_eq!( + declared.get_type(&ctx).deref(&ctx).result_type(), + shared_ptr, + "a declaration keeps its signature" + ); + assert_eq!(call.callee_type(&ctx), declared_type.into()); + module.get_operation().deref(&ctx).verify(&ctx).unwrap(); + } + + #[test] + fn pointer_return_legalization_rejects_an_address_taken_function_before_mutation() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "address_taken".try_into().unwrap()); + let (recover, _) = shared_pointer_recover_function(&mut ctx, &module); + let (call, _) = shared_pointer_caller(&mut ctx, &module, &recover); + let shared_ptr: TypeHandle = PointerType::get(&ctx, 3).into(); + let generic_ptr: TypeHandle = PointerType::get(&ctx, 0).into(); + + let (_taker, entry) = function(&mut ctx, &module, "take_address", generic_ptr, vec![]); + let address = llvm::AddressOfOp::new(&mut ctx, "recover".try_into().unwrap(), 0); + address.get_operation().insert_at_back(entry, &ctx); + let address_value = address.get_operation().deref(&ctx).get_result(0); + llvm::ReturnOp::new(&mut ctx, Some(address_value)) + .get_operation() + .insert_at_back(entry, &ctx); + + let before = operation_ids(&ctx, module.get_operation()); + let error = + crate::legalize_for_nvvm(&mut ctx, module.get_operation(), NvvmIrDialect::Modern, 120) + .unwrap_err(); + let text = error.disp(&ctx).to_string(); + assert!(text.contains("@recover"), "{text}"); + assert!(text.contains("address is taken"), "{text}"); + + assert_eq!(operation_ids(&ctx, module.get_operation()), before); + assert_eq!( + recover.get_type(&ctx).deref(&ctx).result_type(), + shared_ptr, + "nothing changes once the module is rejected" + ); + assert_eq!(call.deref(&ctx).get_result(0).get_type(&ctx), shared_ptr); + } + #[test] fn fneg_is_an_exact_sign_bit_toggle_and_module_still_verifies() { let mut ctx = Context::new(); diff --git a/cuda-oxide/crates/nvvm-transforms/src/lib.rs b/cuda-oxide/crates/nvvm-transforms/src/lib.rs index b14af8ddce..ba0e8ce56b 100644 --- a/cuda-oxide/crates/nvvm-transforms/src/lib.rs +++ b/cuda-oxide/crates/nvvm-transforms/src/lib.rs @@ -24,9 +24,12 @@ use pliron::{ /// Legalize a lowered LLVM module for the selected NVVM input dialect. /// -/// Legacy LLVM 7 modules receive the complete compatibility pass. Modern -/// modules retain modern operations but still rewrite integer widths that -/// libNVVM does not accept in bit-manipulation intrinsics. +/// Both dialects first have non-generic pointer returns of non-kernel +/// functions rewritten to generic pointers, which libNVVM declares correctly +/// at `.func` boundaries. Legacy LLVM 7 modules then receive the complete +/// compatibility pass. Modern modules retain modern operations but still +/// rewrite integer widths that libNVVM does not accept in bit-manipulation +/// intrinsics. /// /// `capability` is the numeric compute capability of the build target /// (`52`, `86`, `90`, ...). The legacy pass uses it to reject rewrites whose @@ -38,6 +41,7 @@ pub fn legalize_for_nvvm( dialect: NvvmIrDialect, capability: u32, ) -> Result<()> { + legalize::legalize_pointer_returns(ctx, module)?; match dialect { NvvmIrDialect::LegacyLlvm7 => legalize::legalize_for_legacy_nvvm(ctx, module, capability), NvvmIrDialect::Modern => legalize::legalize_nvvm_bit_intrinsics(ctx, module), diff --git a/cuda-oxide/crates/oxide-tile-export/Cargo.toml b/cuda-oxide/crates/oxide-tile-export/Cargo.toml new file mode 100644 index 0000000000..35058dd5bd --- /dev/null +++ b/cuda-oxide/crates/oxide-tile-export/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "oxide-tile-export" +version = "0.2.1" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "Descriptions of Rust device functions exported to cuTile tile libraries" +readme = "README.md" + +[dependencies] +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } + +[lints.rust] +missing_docs = { level = "deny" } diff --git a/cuda-oxide/crates/oxide-tile-export/README.md b/cuda-oxide/crates/oxide-tile-export/README.md new file mode 100644 index 0000000000..b94760e86d --- /dev/null +++ b/cuda-oxide/crates/oxide-tile-export/README.md @@ -0,0 +1,21 @@ +# oxide-tile-export + +`oxide-tile-export` defines the description of a Rust device function that a +cuTile program calls through a tile library. + +Three parts of cuda-oxide exchange these descriptions: + +- `#[tile_export]` in `cuda-macros` builds one `TileExportDesc` per exported + function at macro-expansion time and embeds its JSON in a host-only + metadata function. +- The `rustc-codegen-cuda` backend reads the JSON back out of that + function's MIR, validates it, compiles the exported wrapper as a device + root, and writes every description of a crate to + `.tileexports.json` next to the device artifact. +- `cargo oxide` turns that file into the Tile IR interface and implementation + modules of a tile library. + +The crate owns the vocabulary (element types, access modes, synchronization +levels), the validation rules, and the mapping from a description to the +native SIMT ABI that the exported function must have. It performs no file IO +and knows nothing about rustc or MIR. diff --git a/cuda-oxide/crates/oxide-tile-export/src/lib.rs b/cuda-oxide/crates/oxide-tile-export/src/lib.rs new file mode 100644 index 0000000000..dea32f95c1 --- /dev/null +++ b/cuda-oxide/crates/oxide-tile-export/src/lib.rs @@ -0,0 +1,1068 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Descriptions of Rust device functions exported to cuTile tile libraries. +//! +//! `#[tile_export]` (in `cuda-macros`) serializes one [`TileExportDesc`] per +//! exported function at macro-expansion time and embeds it as the body of a +//! host-only metadata function. The codegen backend reads that JSON back out +//! of MIR, validates it, and writes every description of a crate to +//! `.tileexports.json` next to the device artifact. `cargo oxide` +//! turns the file into the Tile IR interface and implementation modules a +//! tile library needs. All three sides share this crate so the vocabulary and +//! the rules are defined once. +//! +//! # The native SIMT ABI +//! +//! A tile library implementation is a plain `void` function. Each declaration +//! argument becomes one parameter, in order, and the result of an exported +//! Rust function becomes a trailing `out` parameter (destination-passing +//! style). The implementation symbol is what the Tile IR call resolves to, so +//! the Rust wrapper must have this signature and no other: +//! +//! | Rust parameter | access | footprint | ABI parameter | +//! |-----------------------|---------|-----------|--------------------------| +//! | `x: T` | `in` | 1 | `T` by value | +//! | `x: &T` | `in` | 1 | `*const T` | +//! | `x: &[T; F]` | `in` | F | `*const T` | +//! | `x: &mut T` | `inout` | 1 | `*mut T` | +//! | `x: &mut [T; F]` | `inout` | F | `*mut T` | +//! | `-> T` (return value) | `out` | 1 | `*mut T`, appended last | +//! +//! Only `x: T` travels by value; every other parameter is a pointer to the +//! invocation's elements, and the function returns nothing. +//! [`TileExportDesc::abi_params`] applies these rules and nothing else. +//! +//! Every parameter of one export has the same footprint: the fragmaps of one +//! implementation share their invocation count, so a 4-element window beside +//! a scalar would need 32 invocations for one argument and 128 for the other. +//! With `sync = warp` the shape must be a multiple of `32 * footprint`, so +//! every invocation group fills whole warps. +//! +//! Thread identity inside the function is the hardware `threadIdx.x`. The +//! callee may not allocate shared memory; version 1 of the description has no +//! scratch buffer. + +#![forbid(unsafe_code)] + +use serde::{Deserialize, Serialize}; + +/// The description format version this crate reads and writes. +/// +/// Version 1 describes one-dimensional tiles, by-value or pointer parameters +/// with a per-invocation footprint, and the `none` and `warp` synchronization +/// levels. [`TileExportDesc::validate`] rejects any other version. +pub const CURRENT_VERSION: u32 = 1; + +/// File suffix of the per-crate description file the backend writes next to +/// the device artifact: `.tileexports.json` holds a JSON array of +/// [`TileExportDesc`]. +pub const DESCRIPTIONS_SUFFIX: &str = "tileexports.json"; + +/// Element type of a tile, named by its Rust type. +/// +/// Tile IR has no unsigned integer types, so the unsigned variants map to the +/// signless `i`-types of the same width in [`TileElem::tile_ir_name`] and to +/// the same LLVM integer types in [`TileElem::llvm_name`]. The Rust name is +/// kept in the description so the wrapper signature and the host-side +/// element interpretation stay recoverable. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum TileElem { + /// IEEE half precision, Rust `f16`, LLVM `half`. + F16, + /// Brain floating point. `cuda-device` carries bf16 values as `u16` bit + /// patterns, so a by-value bf16 is LLVM `i16` in the native ABI (see + /// [`TileElem::llvm_name`]). `#[tile_export]` does not accept `bf16` as + /// an element type; the variant is part of the vocabulary for + /// descriptions written by other tools. + BF16, + /// Rust `f32`, LLVM `float`. + F32, + /// Rust `f64`, LLVM `double`. + F64, + /// Rust `i8`. + I8, + /// Rust `i16`. + I16, + /// Rust `i32`. + I32, + /// Rust `i64`. + I64, + /// Rust `u8`; signless `i8` in Tile IR and LLVM. + U8, + /// Rust `u16`; signless `i16` in Tile IR and LLVM. + U16, + /// Rust `u32`; signless `i32` in Tile IR and LLVM. + U32, + /// Rust `u64`; signless `i64` in Tile IR and LLVM. + U64, +} + +impl TileElem { + /// Every element type, in declaration order. + pub const ALL: [TileElem; 12] = [ + TileElem::F16, + TileElem::BF16, + TileElem::F32, + TileElem::F64, + TileElem::I8, + TileElem::I16, + TileElem::I32, + TileElem::I64, + TileElem::U8, + TileElem::U16, + TileElem::U32, + TileElem::U64, + ]; + + /// The element type whose Rust name is `name`, if any. + /// + /// ``` + /// use oxide_tile_export::TileElem; + /// assert_eq!(TileElem::from_rust_name("f32"), Some(TileElem::F32)); + /// assert_eq!(TileElem::from_rust_name("usize"), None); + /// ``` + pub fn from_rust_name(name: &str) -> Option { + Self::ALL.into_iter().find(|elem| elem.rust_name() == name) + } + + /// The Rust name of the element type (`"f16"`, `"u32"`, ...). + pub fn rust_name(self) -> &'static str { + match self { + TileElem::F16 => "f16", + TileElem::BF16 => "bf16", + TileElem::F32 => "f32", + TileElem::F64 => "f64", + TileElem::I8 => "i8", + TileElem::I16 => "i16", + TileElem::I32 => "i32", + TileElem::I64 => "i64", + TileElem::U8 => "u8", + TileElem::U16 => "u16", + TileElem::U32 => "u32", + TileElem::U64 => "u64", + } + } + + /// The Tile IR name: floats keep their name, integers are signless. + pub fn tile_ir_name(self) -> &'static str { + match self { + TileElem::F16 => "f16", + TileElem::BF16 => "bf16", + TileElem::F32 => "f32", + TileElem::F64 => "f64", + TileElem::I8 | TileElem::U8 => "i8", + TileElem::I16 | TileElem::U16 => "i16", + TileElem::I32 | TileElem::U32 => "i32", + TileElem::I64 | TileElem::U64 => "i64", + } + } + + /// The LLVM IR type the element has in the native ABI. + /// + /// `bf16` is `i16`, not `bfloat`: `cuda-device` has no bf16 scalar type + /// and carries bf16 values as `u16` bit patterns, so that is what the + /// wrapper's parameter lowers to. The Tile IR element stays `bf16`, and + /// `cargo oxide tilelib` records the difference as the implementation's + /// `frontendABI` entry `{bf16 = i16}`. + pub fn llvm_name(self) -> &'static str { + match self { + TileElem::F16 => "half", + TileElem::BF16 => "i16", + TileElem::F32 => "float", + TileElem::F64 => "double", + TileElem::I8 | TileElem::U8 => "i8", + TileElem::I16 | TileElem::U16 => "i16", + TileElem::I32 | TileElem::U32 => "i32", + TileElem::I64 | TileElem::U64 => "i64", + } + } + + /// Width of one element in bits. + pub fn bit_width(self) -> u32 { + match self { + TileElem::I8 | TileElem::U8 => 8, + TileElem::F16 | TileElem::BF16 | TileElem::I16 | TileElem::U16 => 16, + TileElem::F32 | TileElem::I32 | TileElem::U32 => 32, + TileElem::F64 | TileElem::I64 | TileElem::U64 => 64, + } + } + + /// True for the floating-point element types. + pub fn is_float(self) -> bool { + matches!( + self, + TileElem::F16 | TileElem::BF16 | TileElem::F32 | TileElem::F64 + ) + } +} + +/// How the implementation touches one declaration argument. +/// +/// The declaration's results are the tile types of its `out` and `inout` +/// arguments, in order, so every export needs at least one of those. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum TileAccess { + /// Read only. + In, + /// Written, never read; the return value of a Rust function. + Out, + /// Read and written in place. + InOut, +} + +impl TileAccess { + /// The Tile IR name of the access mode (`in`, `out`, `inout`). + pub fn tile_ir_name(self) -> &'static str { + match self { + TileAccess::In => "in", + TileAccess::Out => "out", + TileAccess::InOut => "inout", + } + } + + /// True when the parameter is a result of the declaration. + pub fn is_output(self) -> bool { + matches!(self, TileAccess::Out | TileAccess::InOut) + } +} + +/// How many threads must call the implementation together. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum TileSync { + /// Every invocation is independent. + #[default] + None, + /// Whole warps participate; the shape must be a multiple of 32 times the + /// footprint. + Warp, +} + +impl TileSync { + /// The Tile IR name of the synchronization level. + pub fn tile_ir_name(self) -> &'static str { + match self { + TileSync::None => "none", + TileSync::Warp => "warp", + } + } +} + +/// One parameter of an exported function. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TileParamDesc { + /// Parameter name; the Rust parameter name, or `_out` for the return + /// value. + pub name: String, + /// Element type of the tile this parameter is a fragment of. + pub elem: TileElem, + /// Read, write, or both. + pub access: TileAccess, + /// Number of contiguous elements one invocation owns; at least 1, and the + /// same for every parameter of the export. + pub footprint: u32, + /// True when the element travels by value instead of behind a pointer. + /// Only legal for `in` parameters with footprint 1. + pub by_value: bool, +} + +/// A `#[tile_export]` function as the tile library sees it. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TileExportDesc { + /// Format version; see [`CURRENT_VERSION`]. + pub version: u32, + /// Declaration name: what the cuTile caller names in `foreign_call`. + pub decl: String, + /// Implementation symbol: the `extern "C"` device function the call + /// resolves to. + pub symbol: String, + /// Static tile extents; one positive entry in version 1. + pub shape: Vec, + /// Declaration arguments, in ABI order. + pub params: Vec, + /// Synchronization the implementation requires. + #[serde(default)] + pub sync: TileSync, + /// Compute capabilities the implementation supports (`sm_120`, ...). + /// Empty means every architecture. + #[serde(default)] + pub arch: Vec, +} + +/// What a description can get wrong. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum TileExportError { + /// The description was written by a newer or older format. + #[error( + "tile export description version {found} is not supported; this build reads version {CURRENT_VERSION}" + )] + UnsupportedVersion { + /// The version found in the description. + found: u32, + }, + /// A name that must be an identifier is not one. + #[error( + "{field} `{value}` is not a valid identifier (ASCII letters, digits and `_`, not starting with a digit)" + )] + InvalidIdentifier { + /// Which name: `decl`, `symbol`, or `parameter name`. + field: &'static str, + /// The offending text. + value: String, + }, + /// The shape has the wrong number of dimensions. + #[error("shape must have one dimension in version {CURRENT_VERSION}, found {0}")] + ShapeRank(usize), + /// A shape extent is zero or negative. + #[error("shape extents must be positive, found {0}")] + ShapeNotPositive(i64), + /// The export has no parameters at all. + #[error("a tile export needs at least one parameter")] + NoParams, + /// Nothing is written: no `out` and no `inout` parameter. + #[error("a tile export needs at least one output: return a value or take an `inout` parameter")] + NoOutputs, + /// An `in` parameter sits between two written parameters. The Tile + /// compiler requires the `out` and `inout` arguments of a declaration to + /// be one contiguous run (PreviewOps.cpp, "out/inout fragment args must + /// form a single ..."). + #[error( + "`out` and `inout` parameters must be one contiguous run, but `{between}` is `in` and sits between `{first}` and `{last}`; move `in` parameters before the `&mut` ones" + )] + OutputsNotContiguous { + /// The first written parameter. + first: String, + /// The `in` parameter inside the run. + between: String, + /// The last written parameter (the appended return slot, when present). + last: String, + }, + /// Two parameters share a name. + #[error("parameter `{name}` is declared twice")] + DuplicateParam { + /// The repeated name. + name: String, + }, + /// A footprint of zero elements. + #[error( + "parameter `{name}` has footprint 0; a footprint counts elements and must be at least 1" + )] + ZeroFootprint { + /// The parameter. + name: String, + }, + /// By-value passing on something that is written. + #[error( + "parameter `{name}` is `{access}` and cannot travel by value; only `in` parameters are passed by value" + )] + ByValueNotIn { + /// The parameter. + name: String, + /// Its access mode. + access: &'static str, + }, + /// By-value passing of more than one element. + #[error( + "parameter `{name}` has footprint {footprint} and cannot travel by value; a footprint above 1 needs pointer passing" + )] + ByValueFootprint { + /// The parameter. + name: String, + /// Its footprint. + footprint: u32, + }, + /// Two parameters own a different number of elements per invocation. The + /// fragmaps of one implementation share their invocation count, so this + /// has no Tile IR form. + #[error( + "parameters `{first}` (footprint {first_footprint}) and `{other}` (footprint {other_footprint}) differ; every parameter of one export shares the invocation count, so all footprints must be equal" + )] + MixedFootprints { + /// The first parameter, whose footprint the others were compared to. + first: String, + /// Its footprint. + first_footprint: u32, + /// The first parameter with another footprint. + other: String, + /// That footprint. + other_footprint: u32, + }, + /// Warp synchronization over a shape that does not fill whole warps; the + /// footprint-1 case of [`TileExportError::WarpFootprintShape`]. + #[error("sync = warp needs a shape that is a multiple of 32, found {shape}")] + WarpShape { + /// The offending extent. + shape: i64, + }, + /// Warp synchronization over a shape that does not fill whole warps of + /// footprints: each lane owns `footprint` elements, so a warp owns + /// `32 * footprint` of them. + #[error( + "sync = warp with footprint {footprint} needs a shape that is a multiple of {} (32 lanes times the footprint), found {shape}", + 32 * i64::from(*.footprint) + )] + WarpFootprintShape { + /// The shared footprint. + footprint: u32, + /// The offending extent. + shape: i64, + }, + /// An empty architecture string. + #[error("arch entries must be non-empty names such as `sm_120`")] + EmptyArch, +} + +/// How one ABI parameter is passed. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum AbiParamKind { + /// The element itself, by value. + Scalar(TileElem), + /// A pointer to the invocation's buffer of this element type. + Pointer(TileElem), +} + +impl AbiParamKind { + /// The LLVM IR type of the parameter: the element type, or `ptr`. + pub fn llvm_type(self) -> &'static str { + match self { + AbiParamKind::Scalar(elem) => elem.llvm_name(), + AbiParamKind::Pointer(_) => "ptr", + } + } + + /// The element type behind the parameter. + pub fn elem(self) -> TileElem { + match self { + AbiParamKind::Scalar(elem) | AbiParamKind::Pointer(elem) => elem, + } + } +} + +/// One parameter of the implementation symbol. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AbiParam { + /// The description parameter this ABI parameter carries. + pub name: String, + /// By value or by pointer. + pub kind: AbiParamKind, +} + +/// True for an ASCII identifier: `[A-Za-z_][A-Za-z0-9_]*`. +/// +/// Both Tile IR symbol names and the LLVM symbols the wrapper gets have this +/// form; a Rust raw identifier or a non-ASCII name would not survive the +/// trip. +pub fn is_identifier(text: &str) -> bool { + let mut chars = text.chars(); + match chars.next() { + Some(first) if first.is_ascii_alphabetic() || first == '_' => {} + _ => return false, + } + chars.all(|c| c.is_ascii_alphanumeric() || c == '_') +} + +impl TileExportDesc { + /// Check every rule of version 1. + /// + /// The rules, in the order they are checked: + /// + /// - `version` is [`CURRENT_VERSION`]; + /// - `decl`, `symbol`, and every parameter name are identifiers, and no + /// parameter name repeats; + /// - `shape` has one positive extent; + /// - there is at least one parameter and at least one `out` or `inout` + /// parameter, and the `out`/`inout` parameters are one contiguous run; + /// - every footprint is at least 1, and every parameter has the same + /// footprint; + /// - `by_value` appears only on `in` parameters with footprint 1; + /// - `sync = warp` needs a shape that is a multiple of `32 * footprint`; + /// - every `arch` entry is non-empty. + pub fn validate(&self) -> Result<(), TileExportError> { + if self.version != CURRENT_VERSION { + return Err(TileExportError::UnsupportedVersion { + found: self.version, + }); + } + for (field, value) in [("decl", &self.decl), ("symbol", &self.symbol)] { + if !is_identifier(value) { + return Err(TileExportError::InvalidIdentifier { + field, + value: value.clone(), + }); + } + } + if self.shape.len() != 1 { + return Err(TileExportError::ShapeRank(self.shape.len())); + } + let extent = self.shape[0]; + if extent <= 0 { + return Err(TileExportError::ShapeNotPositive(extent)); + } + if self.params.is_empty() { + return Err(TileExportError::NoParams); + } + let mut seen = std::collections::HashSet::new(); + for param in &self.params { + if !is_identifier(¶m.name) { + return Err(TileExportError::InvalidIdentifier { + field: "parameter name", + value: param.name.clone(), + }); + } + if !seen.insert(param.name.as_str()) { + return Err(TileExportError::DuplicateParam { + name: param.name.clone(), + }); + } + if param.footprint == 0 { + return Err(TileExportError::ZeroFootprint { + name: param.name.clone(), + }); + } + if param.by_value { + if param.access != TileAccess::In { + return Err(TileExportError::ByValueNotIn { + name: param.name.clone(), + access: param.access.tile_ir_name(), + }); + } + if param.footprint != 1 { + return Err(TileExportError::ByValueFootprint { + name: param.name.clone(), + footprint: param.footprint, + }); + } + } + } + // One invocation count per implementation: every fragmap of it maps + // `shape / footprint` invocations, so the footprints cannot differ. + let first = &self.params[0]; + if let Some(other) = self + .params + .iter() + .find(|param| param.footprint != first.footprint) + { + return Err(TileExportError::MixedFootprints { + first: first.name.clone(), + first_footprint: first.footprint, + other: other.name.clone(), + other_footprint: other.footprint, + }); + } + let footprint = first.footprint; + if !self.params.iter().any(|param| param.access.is_output()) { + return Err(TileExportError::NoOutputs); + } + // The written arguments must be adjacent: the Tile compiler ties a + // declaration's results to one contiguous run of out/inout arguments. + let written: Vec = self + .params + .iter() + .enumerate() + .filter(|(_, param)| param.access.is_output()) + .map(|(index, _)| index) + .collect(); + if let (Some(&first), Some(&last)) = (written.first(), written.last()) + && last - first + 1 != written.len() + { + let between = self.params[first + 1..last] + .iter() + .find(|param| !param.access.is_output()) + .expect("a gap in the written run holds an `in` parameter"); + return Err(TileExportError::OutputsNotContiguous { + first: self.params[first].name.clone(), + between: between.name.clone(), + last: self.params[last].name.clone(), + }); + } + // A warp is one invocation of 32 lanes, each owning `footprint` + // elements, so the tile must split into whole `32 * footprint` runs. + if self.sync == TileSync::Warp && extent % (32 * i64::from(footprint)) != 0 { + return Err(if footprint == 1 { + TileExportError::WarpShape { shape: extent } + } else { + TileExportError::WarpFootprintShape { + footprint, + shape: extent, + } + }); + } + if self.arch.iter().any(|arch| arch.trim().is_empty()) { + return Err(TileExportError::EmptyArch); + } + Ok(()) + } + + /// The implementation symbol's parameters, in order. + /// + /// Follows the ABI table in the crate documentation: an `in` parameter + /// marked `by_value` is passed as its element type, and everything else + /// (pointer inputs, `inout` buffers, and the trailing `out` result) is a + /// pointer to the element type. The description is assumed to have passed + /// [`TileExportDesc::validate`]. + pub fn abi_params(&self) -> Vec { + self.params + .iter() + .map(|param| { + let kind = if param.access == TileAccess::In && param.by_value { + AbiParamKind::Scalar(param.elem) + } else { + AbiParamKind::Pointer(param.elem) + }; + AbiParam { + name: param.name.clone(), + kind, + } + }) + .collect() + } + + /// The description as compact JSON, the form the macro embeds. + pub fn to_json(&self) -> String { + serde_json::to_string(self).expect("a tile export description always serializes") + } + + /// The description as indented JSON. + pub fn to_json_pretty(&self) -> String { + serde_json::to_string_pretty(self).expect("a tile export description always serializes") + } + + /// Parse one description. Syntax only; call [`TileExportDesc::validate`] + /// for the semantic rules. + pub fn from_json(text: &str) -> Result { + serde_json::from_str(text) + } +} + +/// Serialize a crate's descriptions as the indented JSON array the +/// description file holds. +pub fn descriptions_to_json_pretty(descs: &[TileExportDesc]) -> String { + serde_json::to_string_pretty(descs).expect("tile export descriptions always serialize") +} + +/// Parse a description file's JSON array. +pub fn descriptions_from_json(text: &str) -> Result, serde_json::Error> { + serde_json::from_str(text) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn param(name: &str, access: TileAccess, footprint: u32, by_value: bool) -> TileParamDesc { + TileParamDesc { + name: name.to_string(), + elem: TileElem::F32, + access, + footprint, + by_value, + } + } + + /// `fn relu(x: f32) -> f32` with `shape = 128`. + fn relu() -> TileExportDesc { + TileExportDesc { + version: CURRENT_VERSION, + decl: "relu".into(), + symbol: "relu_impl".into(), + shape: vec![128], + params: vec![ + param("x", TileAccess::In, 1, true), + param("relu_out", TileAccess::Out, 1, false), + ], + sync: TileSync::None, + arch: Vec::new(), + } + } + + /// `fn scale(x: &mut f32, factor: f32)` exported as `scale_in_place`. + fn scale_in_place() -> TileExportDesc { + TileExportDesc { + version: CURRENT_VERSION, + decl: "scale_in_place".into(), + symbol: "scale_in_place_impl".into(), + shape: vec![128], + params: vec![ + param("x", TileAccess::InOut, 1, false), + param("factor", TileAccess::In, 1, true), + ], + sync: TileSync::None, + arch: vec!["sm_120".into()], + } + } + + #[test] + fn the_examples_validate() { + relu().validate().unwrap(); + scale_in_place().validate().unwrap(); + let mut warp = relu(); + warp.sync = TileSync::Warp; + warp.validate().unwrap(); + } + + #[test] + fn version_is_pinned() { + assert_eq!(CURRENT_VERSION, 1); + let mut desc = relu(); + desc.version = 2; + assert_eq!( + desc.validate(), + Err(TileExportError::UnsupportedVersion { found: 2 }) + ); + } + + #[test] + fn names_must_be_identifiers() { + for (field, mutate) in [ + ( + "decl", + (|d: &mut TileExportDesc| d.decl = "re-lu".into()) as fn(&mut TileExportDesc), + ), + ("symbol", |d: &mut TileExportDesc| d.symbol = "1impl".into()), + ("parameter name", |d: &mut TileExportDesc| { + d.params[0].name = "x y".into() + }), + ] { + let mut desc = relu(); + mutate(&mut desc); + match desc.validate() { + Err(TileExportError::InvalidIdentifier { field: found, .. }) => { + assert_eq!(found, field) + } + other => panic!("{field}: expected an identifier error, got {other:?}"), + } + } + assert!(is_identifier("_x1")); + assert!(!is_identifier("")); + assert!(!is_identifier("9x")); + assert!(!is_identifier("r#type")); + } + + #[test] + fn duplicate_parameter_names_are_rejected() { + let mut desc = relu(); + desc.params[1].name = "x".into(); + assert_eq!( + desc.validate(), + Err(TileExportError::DuplicateParam { name: "x".into() }) + ); + } + + #[test] + fn shape_is_one_positive_extent() { + let mut desc = relu(); + desc.shape = vec![128, 1]; + assert_eq!(desc.validate(), Err(TileExportError::ShapeRank(2))); + desc.shape = Vec::new(); + assert_eq!(desc.validate(), Err(TileExportError::ShapeRank(0))); + desc.shape = vec![0]; + assert_eq!(desc.validate(), Err(TileExportError::ShapeNotPositive(0))); + desc.shape = vec![-4]; + assert_eq!(desc.validate(), Err(TileExportError::ShapeNotPositive(-4))); + } + + #[test] + fn every_export_has_a_parameter_and_an_output() { + let mut desc = relu(); + desc.params.clear(); + assert_eq!(desc.validate(), Err(TileExportError::NoParams)); + let mut desc = relu(); + desc.params.pop(); + assert_eq!(desc.validate(), Err(TileExportError::NoOutputs)); + } + + #[test] + fn written_parameters_form_one_contiguous_run() { + // `fn f(a: &mut f32, b: f32) -> f32` is [inout, in, out]: the Tile + // compiler ties results to one adjacent run of written arguments. + let mut desc = relu(); + desc.params = vec![ + param("a", TileAccess::InOut, 1, false), + param("b", TileAccess::In, 1, true), + param("f_out", TileAccess::Out, 1, false), + ]; + assert_eq!( + desc.validate(), + Err(TileExportError::OutputsNotContiguous { + first: "a".into(), + between: "b".into(), + last: "f_out".into(), + }) + ); + // Grouping the written parameters last is accepted. + desc.params = vec![ + param("b", TileAccess::In, 1, true), + param("a", TileAccess::InOut, 1, false), + param("f_out", TileAccess::Out, 1, false), + ]; + assert_eq!(desc.validate(), Ok(())); + // Written parameters first and `in` parameters after them are also + // one run. + desc.params = vec![ + param("a", TileAccess::InOut, 1, false), + param("c", TileAccess::InOut, 1, false), + param("b", TileAccess::In, 1, true), + ]; + assert_eq!(desc.validate(), Ok(())); + } + + #[test] + fn footprints_start_at_one() { + let mut desc = scale_in_place(); + desc.params[0].footprint = 0; + assert_eq!( + desc.validate(), + Err(TileExportError::ZeroFootprint { name: "x".into() }) + ); + // `fn scale4(x: &mut [f32; 4], factor: &[f32; 4])`. + for param in &mut desc.params { + param.footprint = 4; + param.by_value = false; + } + desc.validate().unwrap(); + } + + #[test] + fn footprints_are_the_same_for_every_parameter() { + // `fn window(acc: &mut [i32; 4], bias: &i32)`: the window needs 16 + // invocations of a 64-element tile, the scalar 64. + let mut desc = scale_in_place(); + desc.params[0].footprint = 4; + assert_eq!( + desc.validate(), + Err(TileExportError::MixedFootprints { + first: "x".into(), + first_footprint: 4, + other: "factor".into(), + other_footprint: 1, + }) + ); + let message = desc.validate().unwrap_err().to_string(); + assert!( + message.contains("`x` (footprint 4) and `factor` (footprint 1) differ"), + "{message}" + ); + // The comparison is against the first parameter, whichever it is. + desc.params[0].footprint = 1; + desc.params[1].footprint = 2; + desc.params[1].by_value = false; + assert_eq!( + desc.validate(), + Err(TileExportError::MixedFootprints { + first: "x".into(), + first_footprint: 1, + other: "factor".into(), + other_footprint: 2, + }) + ); + } + + #[test] + fn by_value_needs_in_and_footprint_one() { + let mut desc = scale_in_place(); + desc.params[0].by_value = true; + assert_eq!( + desc.validate(), + Err(TileExportError::ByValueNotIn { + name: "x".into(), + access: "inout", + }) + ); + let mut desc = relu(); + desc.params[1].by_value = true; + assert_eq!( + desc.validate(), + Err(TileExportError::ByValueNotIn { + name: "relu_out".into(), + access: "out", + }) + ); + let mut desc = relu(); + desc.params[0].footprint = 2; + assert_eq!( + desc.validate(), + Err(TileExportError::ByValueFootprint { + name: "x".into(), + footprint: 2, + }) + ); + } + + #[test] + fn warp_sync_needs_whole_warps() { + let mut desc = relu(); + desc.sync = TileSync::Warp; + desc.shape = vec![100]; + assert_eq!( + desc.validate(), + Err(TileExportError::WarpShape { shape: 100 }) + ); + desc.shape = vec![32]; + desc.validate().unwrap(); + // `none` has no shape rule. + desc.sync = TileSync::None; + desc.shape = vec![100]; + desc.validate().unwrap(); + } + + #[test] + fn warp_sync_with_a_footprint_needs_whole_warps_of_footprints() { + // `fn scale4(x: &mut [f32; 4], factor: &[f32; 4])` with `sync = + // warp`: a warp owns 128 elements, so a 64-element tile holds half + // a warp even though 64 is a multiple of 32. + let mut desc = scale_in_place(); + desc.sync = TileSync::Warp; + for param in &mut desc.params { + param.footprint = 4; + param.by_value = false; + } + desc.shape = vec![64]; + assert_eq!( + desc.validate(), + Err(TileExportError::WarpFootprintShape { + footprint: 4, + shape: 64, + }) + ); + let message = desc.validate().unwrap_err().to_string(); + assert!(message.contains("multiple of 128"), "{message}"); + desc.shape = vec![256]; + desc.validate().unwrap(); + // Footprint 1 keeps the plain multiple-of-32 rule and its message. + for param in &mut desc.params { + param.footprint = 1; + } + desc.shape = vec![100]; + assert_eq!( + desc.validate(), + Err(TileExportError::WarpShape { shape: 100 }) + ); + // Without synchronization only the footprint rules apply. + desc.sync = TileSync::None; + for param in &mut desc.params { + param.footprint = 4; + } + desc.shape = vec![64]; + desc.validate().unwrap(); + } + + #[test] + fn arch_entries_are_non_empty() { + let mut desc = relu(); + desc.arch = vec!["sm_120".into(), " ".into()]; + assert_eq!(desc.validate(), Err(TileExportError::EmptyArch)); + } + + #[test] + fn abi_params_follow_the_table() { + let abi = relu().abi_params(); + assert_eq!( + abi, + vec![ + AbiParam { + name: "x".into(), + kind: AbiParamKind::Scalar(TileElem::F32), + }, + AbiParam { + name: "relu_out".into(), + kind: AbiParamKind::Pointer(TileElem::F32), + }, + ] + ); + assert_eq!(abi[0].kind.llvm_type(), "float"); + assert_eq!(abi[1].kind.llvm_type(), "ptr"); + + let abi = scale_in_place().abi_params(); + assert_eq!(abi[0].kind, AbiParamKind::Pointer(TileElem::F32)); + assert_eq!(abi[1].kind, AbiParamKind::Scalar(TileElem::F32)); + + // `x: &[T; F]` is `in` behind a pointer, `x: &mut [T; F]` is `inout` + // behind a pointer; neither is by value. + let mut desc = relu(); + desc.params[0].by_value = false; + desc.params[0].footprint = 8; + desc.params[1].footprint = 8; + desc.params + .insert(1, param("window", TileAccess::InOut, 8, false)); + desc.validate().unwrap(); + let abi = desc.abi_params(); + assert_eq!(abi[0].kind, AbiParamKind::Pointer(TileElem::F32)); + assert_eq!(abi[1].kind, AbiParamKind::Pointer(TileElem::F32)); + assert_eq!(abi[1].name, "window"); + assert_eq!(abi[2].kind, AbiParamKind::Pointer(TileElem::F32)); + + // A by-value bf16 is the `u16` bit pattern cuda-device carries, so + // the ABI parameter is `i16`; the Tile IR element stays `bf16`. + let mut desc = relu(); + desc.params[0].elem = TileElem::BF16; + assert_eq!(desc.abi_params()[0].kind.llvm_type(), "i16"); + assert_eq!(TileElem::BF16.tile_ir_name(), "bf16"); + } + + #[test] + fn element_names_agree_across_rust_tile_ir_and_llvm() { + let expected = [ + (TileElem::F16, "f16", "f16", "half", 16), + (TileElem::BF16, "bf16", "bf16", "i16", 16), + (TileElem::F32, "f32", "f32", "float", 32), + (TileElem::F64, "f64", "f64", "double", 64), + (TileElem::I8, "i8", "i8", "i8", 8), + (TileElem::I16, "i16", "i16", "i16", 16), + (TileElem::I32, "i32", "i32", "i32", 32), + (TileElem::I64, "i64", "i64", "i64", 64), + (TileElem::U8, "u8", "i8", "i8", 8), + (TileElem::U16, "u16", "i16", "i16", 16), + (TileElem::U32, "u32", "i32", "i32", 32), + (TileElem::U64, "u64", "i64", "i64", 64), + ]; + assert_eq!(expected.len(), TileElem::ALL.len()); + for (elem, rust, tile_ir, llvm, bits) in expected { + assert_eq!(elem.rust_name(), rust); + assert_eq!(elem.tile_ir_name(), tile_ir); + assert_eq!(elem.llvm_name(), llvm); + assert_eq!(elem.bit_width(), bits); + assert_eq!(TileElem::from_rust_name(rust), Some(elem)); + assert_eq!( + elem.is_float(), + matches!(tile_ir, "f16" | "bf16" | "f32" | "f64") + ); + } + assert_eq!(TileElem::from_rust_name("usize"), None); + assert_eq!(TileElem::from_rust_name("bool"), None); + } + + #[test] + fn json_round_trips_and_uses_lowercase_vocabulary() { + let desc = scale_in_place(); + let json = desc.to_json(); + assert!(json.contains("\"access\":\"inout\""), "{json}"); + assert!(json.contains("\"elem\":\"f32\""), "{json}"); + assert!(json.contains("\"sync\":\"none\""), "{json}"); + assert!(json.contains("\"arch\":[\"sm_120\"]"), "{json}"); + assert_eq!(TileExportDesc::from_json(&json).unwrap(), desc); + + let mut warp = relu(); + warp.sync = TileSync::Warp; + warp.params[0].elem = TileElem::BF16; + let json = warp.to_json(); + assert!(json.contains("\"sync\":\"warp\""), "{json}"); + assert!(json.contains("\"elem\":\"bf16\""), "{json}"); + assert_eq!(TileExportDesc::from_json(&json).unwrap(), warp); + + // `sync` and `arch` default when absent. + let minimal = r#"{"version":1,"decl":"relu","symbol":"relu_impl","shape":[128], + "params":[{"name":"x","elem":"f32","access":"in","footprint":1,"by_value":true}, + {"name":"relu_out","elem":"f32","access":"out","footprint":1,"by_value":false}]}"#; + assert_eq!(TileExportDesc::from_json(minimal).unwrap(), relu()); + + let all = vec![relu(), scale_in_place()]; + let text = descriptions_to_json_pretty(&all); + assert!(text.starts_with('['), "{text}"); + assert_eq!(descriptions_from_json(&text).unwrap(), all); + } +} diff --git a/cuda-oxide/crates/oxide-tilebc/Cargo.toml b/cuda-oxide/crates/oxide-tilebc/Cargo.toml new file mode 100644 index 0000000000..6b2cd1fc43 --- /dev/null +++ b/cuda-oxide/crates/oxide-tilebc/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "oxide-tilebc" +version = "0.2.1" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "Tile IR bytecode writer for the tile library descriptions cuda-oxide ships" +readme = "README.md" + +[dependencies] +thiserror = { workspace = true } + +[lints.rust] +missing_docs = { level = "deny" } diff --git a/cuda-oxide/crates/oxide-tilebc/README.md b/cuda-oxide/crates/oxide-tilebc/README.md new file mode 100644 index 0000000000..ecc78100b0 --- /dev/null +++ b/cuda-oxide/crates/oxide-tilebc/README.md @@ -0,0 +1,56 @@ +# oxide-tilebc + +Tile IR bytecode for the tile libraries cuda-oxide ships. + +A tile library lets a Tile kernel call a SIMT function. cuda-oxide is the +SIMT provider: it compiles the Rust function to LTOIR and describes it twice, +once as the declaration a Tile kernel calls (the interface) and once as the +native symbol with the fragment each parameter receives (the implementation). +`tileiras --preview-create-tile-lib` packs the three into one `.tilelib`. This +crate writes the two descriptions; nothing else in the Tile IR format. + +```text +#[tile_export] descriptions ──> Library { declarations, implementations } + │ + interface_bytecode(v) ───┼──> .iface.tilebc (preview$func extern ...) + implementation_bytecode(v)┴──> .impl.tilebc (preview$extern_impl ...) + │ + tileiras --preview-create-tile-lib + .ltoir ──> .tilelib +``` + +The descriptions are public-preview content, which exists from Tile IR 13.5 +on and which an assembler reads only at the version it calls current, the +newest numbered version its `--list-versions` prints. `BytecodeVersion::parse` +reads that listing; pass the newest version to the two writers. + +```rust +use oxide_tilebc::*; + +let tile = TileType::new(vec![128], ScalarType::F32); +let decl = Declaration::new("relu", vec![ + TileArg::new("x", tile.clone(), Access::In), + TileArg::new("relu_out", tile, Access::Out), +]); +let imp = Implementation::new("relu_impl", "relu", vec![ + FragmentArg::new("x", FragmentType::new(vec![1], ScalarType::F32, MemorySpace::Local, + Layout::scalar_mem(), Layout::elementwise(128)), Access::In), + FragmentArg::new("relu_out", FragmentType::new(vec![1], ScalarType::F32, MemorySpace::Local, + Layout::contiguous_mem(1), Layout::elementwise(128)), Access::Out), +]); +let library = Library::new("relu", vec![decl], vec![imp]); +library.verify()?; +let interface = library.interface_bytecode(BytecodeVersion::new(13, 5))?; +let implementation = library.implementation_bytecode(BytecodeVersion::new(13, 5))?; +println!("{}{}", library.interface_text(), library.implementation_text()); +# Ok::<(), oxide_tilebc::Error>(()) +``` + +`verify` applies the rules the Tile compiler's own verifier applies when the +two halves meet: one contiguous run of written arguments, one fragment per +argument with the same access and element type, fragmaps that cover the tile +with the participants the synchronization clause promises, and a rank-0 +buffer layout (a scalar passed by value) only for a read-only local fragment. + +The wire format was derived from the Tile IR compiler sources and checked +against an assembler that reads 13.5 content: its `--preview-create-tile-lib` +packs this crate's output into a library that a Tile C++ kernel then calls. diff --git a/cuda-oxide/crates/oxide-tilebc/src/encoding.rs b/cuda-oxide/crates/oxide-tilebc/src/encoding.rs new file mode 100644 index 0000000000..50816429bf --- /dev/null +++ b/cuda-oxide/crates/oxide-tilebc/src/encoding.rs @@ -0,0 +1,201 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The byte-level primitives of Tile IR bytecode: LEB128 varints, zigzag +//! signed varints, little-endian words, alignment padding and section +//! framing. These are the same primitives every Tile IR producer uses; the +//! reference reader documents them in its bytecode grammar. + +/// The eight magic bytes every Tile IR bytecode file starts with. +pub const MAGIC: [u8; 8] = [0x7F, b'T', b'i', b'l', b'e', b'I', b'R', 0x00]; + +/// Padding byte used to reach a section's alignment. +pub const ALIGNMENT_BYTE: u8 = 0xCB; + +/// Header tag bit set when a module carries public-preview content. +pub const CONTAINS_PREVIEW: u16 = 0x1; + +/// Section identifiers. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub enum Section { + /// Marks the end of the file. + EndOfBytecode = 0x00, + /// The string table. + String = 0x01, + /// The function table. + Func = 0x02, + /// Debug locations and attributes. + Debug = 0x03, + /// The type table. + Type = 0x05, +} + +/// An append-only byte buffer with the bytecode's integer encodings. +#[derive(Debug, Default)] +pub struct Writer { + buf: Vec, +} + +impl Writer { + /// An empty buffer. + pub fn new() -> Self { + Self::default() + } + + /// One raw byte. + pub fn byte(&mut self, value: u8) { + self.buf.push(value); + } + + /// Raw bytes. + pub fn bytes(&mut self, data: &[u8]) { + self.buf.extend_from_slice(data); + } + + /// An unsigned LEB128 varint: seven bits per byte, high bit set while + /// more bytes follow. Tags at or above 0x80 therefore take two bytes. + pub fn varint(&mut self, mut value: u64) { + loop { + let mut byte = (value & 0x7F) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + self.buf.push(byte); + if value == 0 { + return; + } + } + } + + /// A zigzag-encoded signed varint, used for `int64_t` parameters. + pub fn signed_varint(&mut self, value: i64) { + self.varint(((value << 1) ^ (value >> 63)) as u64); + } + + /// A little-endian `u32`. + pub fn le_u32(&mut self, value: u32) { + self.buf.extend_from_slice(&value.to_le_bytes()); + } + + /// A little-endian `u64`. + pub fn le_u64(&mut self, value: u64) { + self.buf.extend_from_slice(&value.to_le_bytes()); + } + + /// `varint(count)` followed by each value as a little-endian `i64`: the + /// encoding of a shape. + pub fn le_var_size_i64(&mut self, values: &[i64]) { + self.varint(values.len() as u64); + for &value in values { + self.buf.extend_from_slice(&value.to_le_bytes()); + } + } + + /// Pad with [`ALIGNMENT_BYTE`] until the length is a multiple of + /// `alignment`. + pub fn align_to(&mut self, alignment: usize) { + if alignment < 2 { + return; + } + while !self.buf.len().is_multiple_of(alignment) { + self.buf.push(ALIGNMENT_BYTE); + } + } + + /// The current length. + #[allow(clippy::len_without_is_empty)] + pub fn len(&self) -> usize { + self.buf.len() + } + + /// Overwrite a little-endian `u32` written earlier (an offset table slot). + pub fn patch_u32(&mut self, offset: usize, value: u32) { + self.buf[offset..offset + 4].copy_from_slice(&value.to_le_bytes()); + } + + /// The bytes. + pub fn into_bytes(self) -> Vec { + self.buf + } +} + +/// Append a section to `out`: the id byte (bit 7 set when aligned), the +/// payload length, then for an aligned section the alignment and padding +/// relative to the whole file, then the payload. +pub fn write_section(out: &mut Vec, section: Section, payload: &[u8], alignment: usize) { + let mut id = section as u8; + if alignment > 1 { + id |= 0x80; + } + let mut header = Writer::new(); + header.byte(id); + header.varint(payload.len() as u64); + if alignment > 1 { + header.varint(alignment as u64); + } + out.extend_from_slice(&header.into_bytes()); + if alignment > 1 { + while !out.len().is_multiple_of(alignment) { + out.push(ALIGNMENT_BYTE); + } + } + out.extend_from_slice(payload); +} + +/// Write `count`, align to four, reserve `count` little-endian `u32` offset +/// slots and return the position of the first one. Tables of strings, types +/// and debug entries all start this way. +pub fn begin_offset_table(w: &mut Writer, count: usize) -> usize { + w.varint(count as u64); + w.align_to(4); + let first = w.len(); + for _ in 0..count { + w.le_u32(0); + } + first +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn varints_follow_leb128() { + let mut w = Writer::new(); + w.varint(0); + w.varint(127); + w.varint(128); + w.varint(0xD2); + w.varint(300); + assert_eq!( + w.into_bytes(), + [0x00, 0x7F, 0x80, 0x01, 0xD2, 0x01, 0xAC, 0x02] + ); + } + + #[test] + fn signed_varints_are_zigzag() { + let mut w = Writer::new(); + w.signed_varint(0); + w.signed_varint(1); + w.signed_varint(-1); + w.signed_varint(128); + assert_eq!(w.into_bytes(), [0x00, 0x02, 0x01, 0x80, 0x02]); + } + + #[test] + fn aligned_sections_pad_relative_to_the_file() { + let mut out = MAGIC.to_vec(); + out.extend_from_slice(&[13, 5, 1, 0]); + write_section(&mut out, Section::Func, &[0xAA; 3], 8); + // id, length, alignment at 12..15, then padding to 16. + assert_eq!(&out[12..16], &[0x82, 0x03, 0x08, ALIGNMENT_BYTE]); + assert_eq!(&out[16..19], &[0xAA; 3]); + write_section(&mut out, Section::EndOfBytecode, &[], 1); + assert_eq!(out.last(), Some(&0x00)); + } +} diff --git a/cuda-oxide/crates/oxide-tilebc/src/lib.rs b/cuda-oxide/crates/oxide-tilebc/src/lib.rs new file mode 100644 index 0000000000..280866cc0f --- /dev/null +++ b/cuda-oxide/crates/oxide-tilebc/src/lib.rs @@ -0,0 +1,360 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Tile IR bytecode for the tile libraries cuda-oxide ships. +//! +//! A tile library lets a Tile kernel call a SIMT function. cuda-oxide is the +//! SIMT provider: it compiles the function to LTOIR and describes it twice, +//! once as the declaration a Tile kernel calls (the interface) and once as +//! the native symbol with the fragment each parameter receives (the +//! implementation). `tileiras --preview-create-tile-lib` packs the three into +//! one `.tilelib`. This crate writes the two descriptions. +//! +//! ```text +//! Library { declarations, implementations } +//! |-- interface_bytecode(version) -> .iface.tilebc +//! |-- implementation_bytecode(version) -> .impl.tilebc +//! `-- interface_text(), implementation_text() readable form +//! ``` +//! +//! The descriptions are public-preview content (Tile IR 13.5 and newer). +//! An assembler reads such content only at the version it calls current, +//! the newest numbered version its `--list-versions` prints, so the caller +//! passes that version; [`BytecodeVersion::parse`] reads the listing. +//! +//! ``` +//! use oxide_tilebc::*; +//! +//! let tile = TileType::new(vec![128], ScalarType::F32); +//! let decl = Declaration::new("relu", vec![ +//! TileArg::new("x", tile.clone(), Access::In), +//! TileArg::new("relu_out", tile, Access::Out), +//! ]); +//! let by_value = FragmentType::new(vec![1], ScalarType::F32, MemorySpace::Local, +//! Layout::scalar_mem(), Layout::elementwise(128)); +//! let by_pointer = FragmentType::new(vec![1], ScalarType::F32, MemorySpace::Local, +//! Layout::contiguous_mem(1), Layout::elementwise(128)); +//! let imp = Implementation::new("relu_impl", "relu", vec![ +//! FragmentArg::new("x", by_value, Access::In), +//! FragmentArg::new("relu_out", by_pointer, Access::Out), +//! ]); +//! let library = Library::new("relu", vec![decl], vec![imp]); +//! library.verify().unwrap(); +//! let interface = library.interface_bytecode(BytecodeVersion::PREVIEW_FIRST).unwrap(); +//! assert_eq!(&interface[8..12], &[13, 5, 1, 0]); +//! ``` + +mod encoding; +mod model; +mod text; +mod verify; +mod writer; + +pub use model::{ + Access, BytecodeVersion, Declaration, FragmentArg, FragmentType, Implementation, Layout, + Library, MemorySpace, ScalarType, Synchronization, TileArg, TileType, +}; + +/// What can go wrong when describing a library. +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum Error { + /// The library breaks a structural rule; the message names the part. + #[error("{0}")] + Verify(String), + /// Preview content needs bytecode 13.5 or newer. + #[error("bytecode version {0} is older than 13.5, the first with public-preview content")] + Version(BytecodeVersion), +} + +#[cfg(test)] +mod tests { + use super::*; + + /// `(section id, payload)` in file order, with the framing checked. + fn sections(bytes: &[u8]) -> Vec<(u8, Vec)> { + assert_eq!(&bytes[..8], b"\x7fTileIR\0"); + let mut pos = 12; + let mut out = Vec::new(); + loop { + let id_byte = bytes[pos]; + pos += 1; + if id_byte & 0x7F == 0 { + assert_eq!(pos, bytes.len()); + return out; + } + let len = read_varint(bytes, &mut pos) as usize; + if id_byte & 0x80 != 0 { + let alignment = read_varint(bytes, &mut pos) as usize; + while pos % alignment != 0 { + assert_eq!(bytes[pos], 0xCB); + pos += 1; + } + } + out.push((id_byte & 0x7F, bytes[pos..pos + len].to_vec())); + pos += len; + } + } + + fn read_varint(bytes: &[u8], pos: &mut usize) -> u64 { + let mut value = 0u64; + let mut shift = 0; + loop { + let byte = bytes[*pos]; + *pos += 1; + value |= u64::from(byte & 0x7F) << shift; + if byte & 0x80 == 0 { + return value; + } + shift += 7; + } + } + + fn section(bytes: &[u8], id: u8) -> Vec { + sections(bytes) + .into_iter() + .find(|(s, _)| *s == id) + .map(|(_, payload)| payload) + .unwrap_or_else(|| panic!("no section {id}")) + } + + fn strings(bytes: &[u8]) -> Vec { + let payload = section(bytes, 1); + let mut pos = 0; + let count = read_varint(&payload, &mut pos) as usize; + while pos % 4 != 0 { + pos += 1; + } + let offsets: Vec = (0..count) + .map(|i| { + u32::from_le_bytes(payload[pos + 4 * i..pos + 4 * i + 4].try_into().unwrap()) + as usize + }) + .collect(); + let data = &payload[pos + 4 * count..]; + (0..count) + .map(|i| { + let end = offsets.get(i + 1).copied().unwrap_or(data.len()); + String::from_utf8(data[offsets[i]..end].to_vec()).unwrap() + }) + .collect() + } + + fn index(table: &[String], text: &str) -> u8 { + u8::try_from(table.iter().position(|t| t == text).expect(text)).unwrap() + } + + fn tile() -> TileType { + TileType::new(vec![128], ScalarType::F32) + } + + fn scale_library() -> Library { + let decl = Declaration::new( + "tilelib_scale", + vec![TileArg::new("arg0", tile(), Access::InOut)], + ); + let fragment = FragmentType::new( + vec![1], + ScalarType::F32, + MemorySpace::Local, + Layout::contiguous_mem(1), + Layout::elementwise(128), + ); + let imp = Implementation::new( + "tilelib_scale_impl", + "tilelib_scale", + vec![FragmentArg::new("arg0", fragment, Access::InOut)], + ); + Library::new("tilelib", vec![decl], vec![imp]) + } + + #[test] + fn header_carries_the_version_and_the_preview_tag() { + let lib = scale_library(); + let bytes = lib.interface_bytecode(BytecodeVersion::new(13, 6)).unwrap(); + assert_eq!(&bytes[8..12], &[13, 6, 1, 0]); + assert_eq!( + lib.interface_bytecode(BytecodeVersion::new(13, 4)), + Err(Error::Version(BytecodeVersion::new(13, 4))) + ); + let ids: Vec = sections(&bytes).into_iter().map(|(id, _)| id).collect(); + assert_eq!(ids, [2, 3, 5, 1], "function, debug, type, string"); + } + + /// Extern declaration: normal form, private, body length 0, then the + /// empty constant parameters. The upstream fixture `tilelib_interface`. + #[test] + fn declaration_entry_layout() { + let bytes = scale_library() + .interface_bytecode(BytecodeVersion::PREVIEW_FIRST) + .unwrap(); + let table = strings(&bytes); + let name = index(&table, "tilelib_scale"); + let access = index(&table, "cuda_tile.preview$access"); + assert_eq!( + section(&bytes, 2), + vec![ + 0x01, name, 0x02, // one function, name, signature type index + 0x01, // flags: private, not a kernel, normal form + 0x01, // debug list index + 0x06, 0x01, 0x0A, 0x01, access, 0xDF, 0x01, + 0x02, // arg_attrs: [{access = inout}] + 0x00, // body length 0: a declaration + 0x06, 0x00, // const_params: empty array + ] + ); + // Types: f32, tile<128xf32>, then (tile) -> (tile). + let types = section(&bytes, 5); + let mut pos = 0; + assert_eq!(read_varint(&types, &mut pos), 3); + let entries = &types[4 + 4 * 3..]; + assert_eq!(entries[0], 0x07, "f32"); + assert_eq!( + &entries[1..12], + &[0x0D, 0x00, 0x01, 128, 0, 0, 0, 0, 0, 0, 0], + "tile<128xf32>" + ); + assert_eq!( + &entries[12..], + &[0x10, 0x01, 0x01, 0x01, 0x01], + "(tile) -> (tile)" + ); + } + + /// Implementation: form bits 01, private, the op payload in place of a + /// body, body length 0. The upstream fixture `tilelib_implementation`. + #[test] + fn implementation_entry_layout() { + let lib = scale_library(); + let bytes = lib + .implementation_bytecode(BytecodeVersion::PREVIEW_FIRST) + .unwrap(); + let table = strings(&bytes); + let symbol = index(&table, "tilelib_scale_impl"); + let implements = index(&table, "tilelib_scale"); + let access = index(&table, "cuda_tile.preview$access"); + let mut expected = vec![ + 0x01, symbol, 0x02, // one function, symbol, signature type index + 0x05, // flags: private, form extern_impl + 0x01, // debug list index + 0x04, // payload flags: arg_attrs only + symbol, implements, // sym_name, implements + 0x02, // function_type index + 0x01, 0x0A, 0x01, access, 0xDF, 0x01, + 0x02, // arg_attrs inline: [{access = inout}] + 0x01, 0x00, // one region, zero blocks + 0x00, // body length 0 + ]; + expected.resize(24, 0xCB); + assert_eq!(section(&bytes, 2), expected); + + // The fragment type entry at 13.5: tagged symbolic shape holding one + // integer per dimension (the 13.5 assembler requires full rank), static shape + // [1], element f32 (index 0), local, (1):(1), (128, 1, 1):(1, 0, 0), + // alignment 1. + let types = section(&bytes, 5); + let entries = &types[4 + 4 * 3..]; + assert_eq!( + &entries[1..], + &[ + 0xD2, 0x01, 0xD8, 0x01, 0x01, 0xD1, 0x01, 0x02, 0x01, 1, 0, 0, 0, 0, 0, 0, 0, 0x00, + 0x00, // + 0xD9, 0x01, 0x01, 0xD1, 0x01, 0x02, 0x01, 0xD1, 0x01, 0x02, // + 0xD9, 0x01, 0x03, 0xD1, 0x01, 0x80, 0x02, 0xD1, 0x01, 0x02, 0xD1, 0x01, + 0x02, // + 0x03, 0xD1, 0x01, 0x02, 0xD1, 0x01, 0x00, 0xD1, 0x01, 0x00, // + 0xD1, 0x01, 0x02, // + 0x10, 0x01, 0x01, 0x00, // (fragment) -> () + ] + ); + } + + /// A warp-synchronous implementation sets bit 8 and writes the value + /// after the argument attributes; arch and frontendABI take bits 6 and 7. + #[test] + fn optional_clauses_set_their_flag_bits() { + let mut lib = scale_library(); + let imp = &mut lib.implementations[0]; + imp.args[0].fragment.fragmap = Layout::warp_lanes(128); + imp.synchronization = Some(Synchronization::Warp); + imp.arch = vec!["sm_120".into()]; + imp.frontend_abi = vec![("bf16".into(), "i16".into())]; + lib.verify().unwrap(); + let bytes = lib + .implementation_bytecode(BytecodeVersion::PREVIEW_FIRST) + .unwrap(); + let table = strings(&bytes); + let func = section(&bytes, 2); + // flags = bits 2, 6, 7, 8 = 0x1C4 -> varint C4 03 + assert_eq!(&func[5..7], &[0xC4, 0x03]); + let after_attrs = 7 + 3 + 7; + assert_eq!( + &func[after_attrs..after_attrs + 10], + &[ + 0x01, + 0x05, + index(&table, "sm_120"), // arch: one tagged string + 0x01, + index(&table, "bf16"), + 0x05, + index(&table, "i16"), // frontendABI + 0x01, // synchronization = warp + 0x01, + 0x00, // region + ] + ); + } + + #[test] + fn debug_section_holds_one_unknown_location_per_function() { + let bytes = scale_library() + .interface_bytecode(BytecodeVersion::PREVIEW_FIRST) + .unwrap(); + assert_eq!( + section(&bytes, 3), + vec![ + 0x01, 0xCB, 0xCB, 0xCB, // one function, aligned + 0, 0, 0, 0, // its id list starts at 0 + 0x01, 0xCB, 0xCB, 0xCB, 0xCB, 0xCB, 0xCB, 0xCB, // one id, aligned to 8 + 0, 0, 0, 0, 0, 0, 0, 0, // id 0: unknown location + 0x01, 0xCB, 0xCB, 0xCB, // one attribute entry + 0, 0, 0, 0, // at offset 0 + 0x00, // tag 0: unknown + ] + ); + } + + #[test] + fn text_matches_the_reference_form() { + let lib = scale_library(); + assert_eq!( + lib.interface_text(), + "cuda_tile.module @tilelib {\n cuda_tile.preview$func extern @tilelib_scale(%arg0: \ + !cuda_tile.tile<128xf32> {cuda_tile.preview$access = #cuda_tile.preview$access}) \ + -> !cuda_tile.tile<128xf32>\n}\n" + ); + assert!(lib.implementation_text().contains( + "@tilelib_scale_impl implements @tilelib_scale(%arg0: !cuda_tile.preview$fragment<1xf32, \ + local, mem_layout = #cuda_tile.preview$layout<(1):(1)>, fragmap = \ + #cuda_tile.preview$layout<(128, 1, 1):(1, 0, 0)>> {cuda_tile.preview$access = \ + #cuda_tile.preview$access})\n" + )); + } + + #[test] + fn versions_parse_from_the_assembler_listing() { + assert_eq!( + BytecodeVersion::parse("13.5"), + Some(BytecodeVersion::new(13, 5)) + ); + assert_eq!( + BytecodeVersion::parse(" 13.6.1 "), + Some(BytecodeVersion::new(13, 6)) + ); + assert_eq!(BytecodeVersion::parse("internal"), None); + assert_eq!(BytecodeVersion::parse("13"), None); + assert!(BytecodeVersion::new(13, 6) > BytecodeVersion::PREVIEW_FIRST); + assert_eq!(BytecodeVersion::new(13, 6).to_string(), "13.6"); + } +} diff --git a/cuda-oxide/crates/oxide-tilebc/src/model.rs b/cuda-oxide/crates/oxide-tilebc/src/model.rs new file mode 100644 index 0000000000..d6ef3991fc --- /dev/null +++ b/cuda-oxide/crates/oxide-tilebc/src/model.rs @@ -0,0 +1,465 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The vocabulary of a tile library description: the tiles a declaration +//! takes, the fragments an implementation receives, and the clauses that +//! tell the Tile compiler how to call the native function. + +use std::fmt; + +/// A Tile IR bytecode version (`major.minor`). The header tag bits are not +/// part of a version. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct BytecodeVersion { + /// Major version, 13 for every Tile IR release so far. + pub major: u8, + /// Minor version. + pub minor: u8, +} + +impl BytecodeVersion { + /// The first version with public-preview content, and the version the + /// descriptions are written at when no assembler can be asked. + pub const PREVIEW_FIRST: Self = Self { + major: 13, + minor: 5, + }; + + /// A version from its two numbers. + pub const fn new(major: u8, minor: u8) -> Self { + Self { major, minor } + } + + /// Parse one line of `tileiras --list-versions`: `13.5` or `13.5.1` + /// (the header tag is ignored). Lines that are not versions, such as a + /// named build an assembler may list, give `None`. + pub fn parse(text: &str) -> Option { + let mut parts = text.trim().split('.'); + let major = parts.next()?.parse().ok()?; + let minor = parts.next()?.parse().ok()?; + if let Some(tag) = parts.next() { + tag.parse::().ok()?; + } + if parts.next().is_some() { + return None; + } + Some(Self { major, minor }) + } +} + +impl fmt::Display for BytecodeVersion { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}.{}", self.major, self.minor) + } +} + +/// Element types a description can name. Tile IR has no unsigned integers: +/// the unsigned Rust types are the signless `iN` of the same width. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ScalarType { + /// `f16` + F16, + /// `bf16` + BF16, + /// `f32` + F32, + /// `f64` + F64, + /// `i8` + I8, + /// `i16` + I16, + /// `i32` + I32, + /// `i64` + I64, +} + +impl ScalarType { + /// The type-table tag. + pub(crate) fn tag(self) -> u8 { + match self { + Self::I8 => 1, + Self::I16 => 2, + Self::I32 => 3, + Self::I64 => 4, + Self::F16 => 5, + Self::BF16 => 6, + Self::F32 => 7, + Self::F64 => 9, + } + } + + /// The textual name. + pub fn name(self) -> &'static str { + match self { + Self::F16 => "f16", + Self::BF16 => "bf16", + Self::F32 => "f32", + Self::F64 => "f64", + Self::I8 => "i8", + Self::I16 => "i16", + Self::I32 => "i32", + Self::I64 => "i64", + } + } +} + +/// A concrete tile type, the value a declaration argument carries. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct TileType { + /// Extents, one per dimension. + pub shape: Vec, + /// Element type. + pub element: ScalarType, +} + +impl TileType { + /// A tile of `shape` by `element`. + pub fn new(shape: Vec, element: ScalarType) -> Self { + Self { shape, element } + } + + /// Number of elements. + pub fn size(&self) -> i64 { + self.shape.iter().product() + } +} + +/// How a callee uses an argument. `out` and `inout` arguments are also the +/// declaration's results, in argument order. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Access { + /// Read only. + In, + /// Written, not read. + Out, + /// Read and written. + InOut, +} + +impl Access { + /// The bytecode value. + pub(crate) fn value(self) -> u64 { + match self { + Self::In => 0, + Self::Out => 1, + Self::InOut => 2, + } + } + + /// The textual name. + pub fn name(self) -> &'static str { + match self { + Self::In => "in", + Self::Out => "out", + Self::InOut => "inout", + } + } + + /// Whether the callee writes the argument. + pub fn is_written(self) -> bool { + !matches!(self, Self::In) + } +} + +/// Where a fragment's buffer lives. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MemorySpace { + /// Per-thread storage. + Local, + /// CTA shared memory. + Smem, +} + +impl MemorySpace { + pub(crate) fn value(self) -> u64 { + match self { + Self::Local => 0, + Self::Smem => 1, + } + } + + /// The textual name. + pub fn name(self) -> &'static str { + match self { + Self::Local => "local", + Self::Smem => "smem", + } + } +} + +/// The threads that call the native function together. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Synchronization { + /// One thread per invocation. + None, + /// The 32 lanes of a warp. + Warp, + /// Four warps. + Warpgroup, + /// The whole CTA. + Cta, +} + +impl Synchronization { + pub(crate) fn value(self) -> u64 { + match self { + Self::None => 0, + Self::Warp => 1, + Self::Warpgroup => 2, + Self::Cta => 3, + } + } + + /// The textual name. + pub fn name(self) -> &'static str { + match self { + Self::None => "none", + Self::Warp => "warp", + Self::Warpgroup => "warpgroup", + Self::Cta => "cta", + } + } + + /// How many threads one invocation spans. + pub fn participants(self) -> i64 { + match self { + Self::None => 1, + Self::Warp => 32, + Self::Warpgroup => 128, + Self::Cta => 0, + } + } +} + +/// A concrete layout: `(shapes):(strides)`, integers only. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Layout { + /// Mode extents. + pub shapes: Vec, + /// Mode strides. + pub strides: Vec, +} + +impl Layout { + /// A layout from its shapes and strides. + pub fn flat(shapes: &[i64], strides: &[i64]) -> Self { + Self { + shapes: shapes.to_vec(), + strides: strides.to_vec(), + } + } + + /// The rank-0 buffer layout `():()` of a scalar passed by value. + pub fn scalar_mem() -> Self { + Self::flat(&[], &[]) + } + + /// The buffer layout `(F):(1)` of `footprint` contiguous elements. + pub fn contiguous_mem(footprint: i64) -> Self { + Self::flat(&[footprint], &[1]) + } + + /// One invocation per element, no participants: `(N, 1, 1):(1, 0, 0)`. + pub fn elementwise(tile_size: i64) -> Self { + Self::flat(&[tile_size, 1, 1], &[1, 0, 0]) + } + + /// Warp-synchronous, one element per lane: `(N/32, 32, 1):(32, 1, 0)`. + pub fn warp_lanes(tile_size: i64) -> Self { + Self::flat(&[tile_size / 32, 32, 1], &[32, 1, 0]) + } + + /// One invocation per block of `footprint` contiguous elements: + /// `(N/F, 1, F):(F, 0, 1)`. + pub fn footprint_blocks(tile_size: i64, footprint: i64) -> Self { + Self::flat(&[tile_size / footprint, 1, footprint], &[footprint, 0, 1]) + } + + /// Number of dimensions. + pub fn rank(&self) -> usize { + self.shapes.len() + } + + /// Product of the shapes. + pub fn size(&self) -> i64 { + self.shapes.iter().product() + } +} + +/// The fragment the native function receives for one declaration argument. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct FragmentType { + /// Elements one invocation owns (its footprint), one extent per dimension. + pub shape: Vec, + /// Element type. + pub element: ScalarType, + /// Where the buffer lives. + pub memory_space: MemorySpace, + /// How the footprint sits in the buffer; rank 0 means passed by value. + pub mem_layout: Layout, + /// `(invocations, participants, footprint):(strides)` onto the tile. + pub fragmap: Layout, +} + +impl FragmentType { + /// A fragment from its parts. The alignment the bytecode carries is + /// always the default of one. + pub fn new( + shape: Vec, + element: ScalarType, + memory_space: MemorySpace, + mem_layout: Layout, + fragmap: Layout, + ) -> Self { + Self { + shape, + element, + memory_space, + mem_layout, + fragmap, + } + } + + /// Elements per invocation. + pub fn footprint(&self) -> i64 { + self.shape.iter().product() + } +} + +/// One argument of a declaration. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TileArg { + /// Argument name, printed as `%name`. + pub name: String, + /// The tile at the call boundary. + pub ty: TileType, + /// How the callee uses it. + pub access: Access, +} + +impl TileArg { + /// An argument from its parts. + pub fn new(name: impl Into, ty: TileType, access: Access) -> Self { + Self { + name: name.into(), + ty, + access, + } + } +} + +/// One argument of an implementation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FragmentArg { + /// Argument name, printed as `%name`. + pub name: String, + /// The fragment the native function receives. + pub fragment: FragmentType, + /// Must equal the access of the matching declaration argument. + pub access: Access, +} + +impl FragmentArg { + /// An argument from its parts. + pub fn new(name: impl Into, fragment: FragmentType, access: Access) -> Self { + Self { + name: name.into(), + fragment, + access, + } + } +} + +/// An extern declaration: what a Tile kernel can call. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Declaration { + /// The symbol a `preview$call` names. + pub name: String, + /// Arguments in order; at least one. + pub args: Vec, +} + +impl Declaration { + /// A declaration from its name and arguments. + pub fn new(name: impl Into, args: Vec) -> Self { + Self { + name: name.into(), + args, + } + } + + /// The tiles the declaration returns: its written arguments, in order. + pub fn results(&self) -> impl Iterator { + self.args + .iter() + .filter(|arg| arg.access.is_written()) + .map(|arg| &arg.ty) + } +} + +/// An implementation: the native symbol the code payload defines and how +/// its parameters map onto a declaration's tiles. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Implementation { + /// The linked symbol. + pub symbol: String, + /// The declaration this implements. + pub implements: String, + /// One fragment per declaration argument, in order. + pub args: Vec, + /// The `synchronization` clause; `None` leaves the clause out, which + /// means one thread per invocation. + pub synchronization: Option, + /// Target names such as `sm_120`; empty matches every target. + pub arch: Vec, + /// `frontendABI` entries: printed type to native type. + pub frontend_abi: Vec<(String, String)>, +} + +impl Implementation { + /// An implementation with no optional clauses. + pub fn new( + symbol: impl Into, + implements: impl Into, + args: Vec, + ) -> Self { + Self { + symbol: symbol.into(), + implements: implements.into(), + args, + synchronization: None, + arch: Vec::new(), + frontend_abi: Vec::new(), + } + } +} + +/// A tile library's two halves, under one module name. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Library { + /// The `cuda_tile.module` name of both halves. + pub name: String, + /// The interface: what can be called. + pub declarations: Vec, + /// The catalog: how each declaration is implemented. + pub implementations: Vec, +} + +impl Library { + /// A library from its parts. + pub fn new( + name: impl Into, + declarations: Vec, + implementations: Vec, + ) -> Self { + Self { + name: name.into(), + declarations, + implementations, + } + } +} diff --git a/cuda-oxide/crates/oxide-tilebc/src/text.rs b/cuda-oxide/crates/oxide-tilebc/src/text.rs new file mode 100644 index 0000000000..32073f2b6e --- /dev/null +++ b/cuda-oxide/crates/oxide-tilebc/src/text.rs @@ -0,0 +1,134 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The readable form of a library: the Tile IR textual syntax of both +//! modules, as the reference tools print it. It is written next to the +//! bytecode so a reader can see what the assembler will read. + +use std::fmt::Write as _; + +use crate::model::{Declaration, FragmentType, Implementation, Layout, Library, TileType}; + +fn layout(layout: &Layout) -> String { + let list = |values: &[i64]| { + values + .iter() + .map(|v| v.to_string()) + .collect::>() + .join(", ") + }; + format!( + "#cuda_tile.preview$layout<({}):({})>", + list(&layout.shapes), + list(&layout.strides) + ) +} + +fn tile(ty: &TileType) -> String { + let dims: Vec = ty.shape.iter().map(|d| d.to_string()).collect(); + format!("!cuda_tile.tile<{}x{}>", dims.join("x"), ty.element.name()) +} + +fn fragment(ty: &FragmentType) -> String { + let dims: Vec = ty.shape.iter().map(|d| d.to_string()).collect(); + format!( + "!cuda_tile.preview$fragment<{}x{}, {}, mem_layout = {}, fragmap = {}>", + dims.join("x"), + ty.element.name(), + ty.memory_space.name(), + layout(&ty.mem_layout), + layout(&ty.fragmap) + ) +} + +fn access(access: crate::model::Access) -> String { + format!( + "{{cuda_tile.preview$access = #cuda_tile.preview$access<{}>}}", + access.name() + ) +} + +fn declaration(decl: &Declaration) -> String { + let args: Vec = decl + .args + .iter() + .map(|arg| format!("%{}: {} {}", arg.name, tile(&arg.ty), access(arg.access))) + .collect(); + let results: Vec = decl.results().map(tile).collect(); + let results = match results.len() { + 1 => results[0].clone(), + _ => format!("({})", results.join(", ")), + }; + format!( + " cuda_tile.preview$func extern @{}({}) -> {}\n", + decl.name, + args.join(", "), + results + ) +} + +fn implementation(imp: &Implementation) -> String { + let args: Vec = imp + .args + .iter() + .map(|arg| { + format!( + "%{}: {} {}", + arg.name, + fragment(&arg.fragment), + access(arg.access) + ) + }) + .collect(); + let mut text = format!( + " cuda_tile.preview$extern_impl @{} implements @{}({})", + imp.symbol, + imp.implements, + args.join(", ") + ); + if let Some(sync) = imp.synchronization { + let _ = write!( + text, + "\n synchronization = #cuda_tile.preview$synchronization<{}>", + sync.name() + ); + } + if !imp.arch.is_empty() { + let names: Vec = imp.arch.iter().map(|name| format!("\"{name}\"")).collect(); + let _ = write!(text, "\n arch = [{}]", names.join(", ")); + } + if !imp.frontend_abi.is_empty() { + let entries: Vec = imp + .frontend_abi + .iter() + .map(|(key, value)| format!("{key} = \"{value}\"")) + .collect(); + let _ = write!(text, "\n {{frontendABI = {{{}}}}}", entries.join(", ")); + } + text.push('\n'); + text +} + +impl Library { + /// The interface module in textual Tile IR. + pub fn interface_text(&self) -> String { + let mut text = format!("cuda_tile.module @{} {{\n", self.name); + for decl in &self.declarations { + text.push_str(&declaration(decl)); + } + text.push_str("}\n"); + text + } + + /// The implementation module in textual Tile IR. + pub fn implementation_text(&self) -> String { + let mut text = format!("cuda_tile.module @{} {{\n", self.name); + for imp in &self.implementations { + text.push_str(&implementation(imp)); + } + text.push_str("}\n"); + text + } +} diff --git a/cuda-oxide/crates/oxide-tilebc/src/verify.rs b/cuda-oxide/crates/oxide-tilebc/src/verify.rs new file mode 100644 index 0000000000..3cc124448b --- /dev/null +++ b/cuda-oxide/crates/oxide-tilebc/src/verify.rs @@ -0,0 +1,360 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The structural rules the Tile compiler's verifier applies when a library's +//! two halves meet, checked here before any bytecode is written so a mistake +//! is reported in terms of the description rather than of a deserialized +//! module. + +use std::collections::HashSet; + +use crate::Error; +use crate::model::{Access, Declaration, FragmentArg, Implementation, Library, MemorySpace}; + +fn is_identifier(text: &str) -> bool { + let mut chars = text.chars(); + matches!(chars.next(), Some(c) if c == '_' || c.is_ascii_alphabetic()) + && chars.all(|c| c == '_' || c.is_ascii_alphanumeric()) +} + +impl Library { + /// Check the library against the rules below; the first violation is + /// the error. + /// + /// - names are identifiers, unique among declarations and among symbols; + /// - a declaration has at least one argument, one dimension per tile, + /// and its written arguments form one contiguous run (the native + /// implementation path needs destination-passing results in order); + /// - every implementation names a declaration, with one fragment per + /// argument, the same access and element type; + /// - a fragment covers the tile: invocations times participants times + /// footprint equals the tile size, the footprint equals the fragment's + /// shape, a rank-0 buffer layout only for a read-only local fragment; + /// - the participants of every fragmap equal what the synchronization + /// clause promises (1 without a clause, 32 for a warp). + pub fn verify(&self) -> Result<(), Error> { + if !is_identifier(&self.name) { + return Err(Error::Verify(format!( + "module name `{}` is not an identifier", + self.name + ))); + } + let mut names = HashSet::new(); + for decl in &self.declarations { + verify_declaration(decl)?; + if !names.insert(decl.name.as_str()) { + return Err(Error::Verify(format!( + "declaration @{} appears twice", + decl.name + ))); + } + } + let mut symbols = HashSet::new(); + for imp in &self.implementations { + if !symbols.insert(imp.symbol.as_str()) { + return Err(Error::Verify(format!( + "implementation @{} appears twice", + imp.symbol + ))); + } + let decl = self + .declarations + .iter() + .find(|decl| decl.name == imp.implements) + .ok_or_else(|| { + Error::Verify(format!( + "implementation @{} implements @{}, which the interface does not declare", + imp.symbol, imp.implements + )) + })?; + verify_implementation(imp, decl)?; + } + Ok(()) + } +} + +fn verify_declaration(decl: &Declaration) -> Result<(), Error> { + let at = |message: String| Error::Verify(format!("declaration @{}: {message}", decl.name)); + if !is_identifier(&decl.name) { + return Err(Error::Verify(format!( + "declaration name `{}` is not an identifier", + decl.name + ))); + } + if decl.args.is_empty() { + return Err(at("a declaration needs at least one argument".into())); + } + let mut seen = HashSet::new(); + for arg in &decl.args { + if !is_identifier(&arg.name) { + return Err(at(format!( + "argument name `{}` is not an identifier", + arg.name + ))); + } + if !seen.insert(arg.name.as_str()) { + return Err(at(format!("argument `{}` appears twice", arg.name))); + } + if arg.ty.shape.len() != 1 { + return Err(at(format!( + "argument `{}` has {} dimensions; this writer emits one-dimensional tiles", + arg.name, + arg.ty.shape.len() + ))); + } + if arg.ty.shape[0] <= 0 { + return Err(at(format!( + "argument `{}` has a non-positive extent", + arg.name + ))); + } + } + let written: Vec = decl + .args + .iter() + .enumerate() + .filter(|(_, arg)| arg.access.is_written()) + .map(|(i, _)| i) + .collect(); + if written.is_empty() { + return Err(at( + "no argument is written; a declaration's results are its out and inout arguments" + .into(), + )); + } + if written.last().unwrap() - written[0] + 1 != written.len() { + return Err(at(format!( + "written arguments must form one contiguous run; `{}` sits between them", + decl.args[written[0] + 1..*written.last().unwrap()] + .iter() + .find(|arg| !arg.access.is_written()) + .map(|arg| arg.name.as_str()) + .unwrap_or("?") + ))); + } + Ok(()) +} + +fn verify_implementation(imp: &Implementation, decl: &Declaration) -> Result<(), Error> { + let at = |message: String| Error::Verify(format!("implementation @{}: {message}", imp.symbol)); + if !is_identifier(&imp.symbol) { + return Err(Error::Verify(format!( + "implementation symbol `{}` is not an identifier", + imp.symbol + ))); + } + if imp.args.len() != decl.args.len() { + return Err(at(format!( + "{} fragment arguments for the {} arguments of @{}", + imp.args.len(), + decl.args.len(), + decl.name + ))); + } + let participants = imp.synchronization.map_or(1, |sync| sync.participants()); + for (arg, decl_arg) in imp.args.iter().zip(&decl.args) { + verify_fragment(arg, decl_arg.ty.size(), decl_arg.ty.element, participants) + .map_err(|message| at(format!("argument `{}`: {message}", arg.name)))?; + if arg.access != decl_arg.access { + return Err(at(format!( + "argument `{}` is `{}` but the declaration says `{}`", + arg.name, + arg.access.name(), + decl_arg.access.name() + ))); + } + } + for name in &imp.arch { + if name.is_empty() { + return Err(at("an empty arch name".into())); + } + } + Ok(()) +} + +fn verify_fragment( + arg: &FragmentArg, + tile_size: i64, + element: crate::model::ScalarType, + participants: i64, +) -> Result<(), String> { + let fragment = &arg.fragment; + if fragment.element != element { + return Err(format!( + "fragment element type {} differs from the tile element type {}", + fragment.element.name(), + element.name() + )); + } + if fragment.shape.is_empty() || fragment.shape.iter().any(|d| *d <= 0) { + return Err("fragment shape must have positive extents".into()); + } + let footprint = fragment.footprint(); + let fragmap = &fragment.fragmap; + if fragmap.rank() != 3 || fragmap.strides.len() != 3 { + return Err(format!( + "fragmap must have three modes (invocations, participants, footprint), found {}", + fragmap.rank() + )); + } + if fragmap.shapes[2] != footprint { + return Err(format!( + "fragmap footprint {} differs from the fragment's {} elements", + fragmap.shapes[2], footprint + )); + } + if fragmap.size() != tile_size { + return Err(format!( + "fragmap covers {} elements of a {tile_size}-element tile", + fragmap.size() + )); + } + if participants != 0 && fragmap.shapes[1] != participants { + return Err(format!( + "fragmap has {} participants but the synchronization clause promises {participants}", + fragmap.shapes[1] + )); + } + let mem = &fragment.mem_layout; + if mem.shapes.len() != mem.strides.len() { + return Err("mem_layout shapes and strides differ in length".into()); + } + if mem.rank() == 0 { + if footprint != 1 { + return Err("a rank-0 mem_layout needs a one-element footprint".into()); + } + if arg.access != Access::In || fragment.memory_space != MemorySpace::Local { + return Err( + "a rank-0 mem_layout (passed by value) is only valid for a read-only local fragment" + .into(), + ); + } + } else if mem.size() != footprint { + return Err(format!( + "mem_layout spans {} elements but the footprint is {footprint}", + mem.size() + )); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::model::*; + + fn tile() -> TileType { + TileType::new(vec![128], ScalarType::F32) + } + + fn scalar_fragment() -> FragmentType { + FragmentType::new( + vec![1], + ScalarType::F32, + MemorySpace::Local, + Layout::scalar_mem(), + Layout::elementwise(128), + ) + } + + fn pointer_fragment() -> FragmentType { + FragmentType::new( + vec![1], + ScalarType::F32, + MemorySpace::Local, + Layout::contiguous_mem(1), + Layout::elementwise(128), + ) + } + + fn relu() -> Library { + Library::new( + "lib", + vec![Declaration::new( + "relu", + vec![ + TileArg::new("x", tile(), Access::In), + TileArg::new("relu_out", tile(), Access::Out), + ], + )], + vec![Implementation::new( + "relu_impl", + "relu", + vec![ + FragmentArg::new("x", scalar_fragment(), Access::In), + FragmentArg::new("relu_out", pointer_fragment(), Access::Out), + ], + )], + ) + } + + #[test] + fn the_relu_library_verifies() { + relu().verify().unwrap(); + } + + #[test] + fn written_arguments_must_be_contiguous() { + let mut lib = relu(); + lib.declarations[0] + .args + .push(TileArg::new("y", tile(), Access::In)); + lib.declarations[0] + .args + .push(TileArg::new("z", tile(), Access::Out)); + let error = lib.verify().unwrap_err().to_string(); + assert!( + error.contains("contiguous run") && error.contains("`y`"), + "{error}" + ); + } + + #[test] + fn an_implementation_needs_its_declaration() { + let mut lib = relu(); + lib.implementations[0].implements = "gelu".into(); + let error = lib.verify().unwrap_err().to_string(); + assert!(error.contains("does not declare"), "{error}"); + } + + #[test] + fn access_and_coverage_are_checked() { + let mut lib = relu(); + lib.implementations[0].args[1].access = Access::InOut; + assert!( + lib.verify() + .unwrap_err() + .to_string() + .contains("`relu_out` is `inout`") + ); + + let mut lib = relu(); + lib.implementations[0].args[1].fragment.fragmap = Layout::elementwise(64); + assert!( + lib.verify() + .unwrap_err() + .to_string() + .contains("covers 64 elements") + ); + + let mut lib = relu(); + lib.implementations[0].synchronization = Some(Synchronization::Warp); + let error = lib.verify().unwrap_err().to_string(); + assert!( + error.contains("1 participants") && error.contains("promises 32"), + "{error}" + ); + lib.implementations[0].args[0].fragment.fragmap = Layout::warp_lanes(128); + lib.implementations[0].args[1].fragment.fragmap = Layout::warp_lanes(128); + lib.verify().unwrap(); + } + + #[test] + fn by_value_is_only_for_read_only_local_scalars() { + let mut lib = relu(); + lib.implementations[0].args[1].fragment.mem_layout = Layout::scalar_mem(); + let error = lib.verify().unwrap_err().to_string(); + assert!(error.contains("rank-0 mem_layout"), "{error}"); + } +} diff --git a/cuda-oxide/crates/oxide-tilebc/src/writer.rs b/cuda-oxide/crates/oxide-tilebc/src/writer.rs new file mode 100644 index 0000000000..75beff3444 --- /dev/null +++ b/cuda-oxide/crates/oxide-tilebc/src/writer.rs @@ -0,0 +1,406 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Serialize a library's two halves as Tile IR bytecode. +//! +//! The two modules only ever contain bodiless functions, so the file has +//! four sections: the function table (one entry per declaration or +//! implementation), the debug table (one unknown location per function), the +//! type table and the string table. Every encoding below was checked against +//! the reference reader's sources and against an assembler that reads 13.5 +//! content, which packs the output into a `.tilelib` with +//! `--preview-create-tile-lib`. +//! +//! ```text +//! header: 7F 'T' 'i' 'l' 'e' 'I' 'R' 00 | major | minor | tag u16 (bit 0: preview) +//! section Func (aligned 8): varint(n) then entries, padded to 8 +//! section Debug (aligned 8): per-function id lists, one unknown entry +//! section Type (aligned 4): varint(n), offsets, entries +//! section String(aligned 4): varint(n), offsets, bytes +//! end byte 00 +//! ``` + +use std::collections::HashMap; + +use crate::Error; +use crate::encoding::{ + CONTAINS_PREVIEW, MAGIC, Section, Writer, begin_offset_table, write_section, +}; +use crate::model::{ + Access, BytecodeVersion, Declaration, FragmentType, Implementation, Layout, Library, + ScalarType, TileType, +}; + +// Type-table tags. +const TAG_TILE: u64 = 13; +const TAG_FUNC: u64 = 16; +const TAG_FRAGMENT: u64 = 0xD2; + +// Attribute tags. +const ATTR_STRING: u64 = 5; +const ATTR_ARRAY: u64 = 6; +const ATTR_DICTIONARY: u64 = 10; +const ATTR_SYMBOLIC_INT: u64 = 0xD1; +const ATTR_SYMBOLIC_TUPLE: u64 = 0xD8; +const ATTR_LAYOUT: u64 = 0xD9; +const ATTR_ACCESS: u64 = 0xDF; + +/// The argument attribute that carries an access. +const ACCESS_ATTR: &str = "cuda_tile.preview$access"; + +// Function-table entry flags: bit 0 private, bit 1 kernel, bits 3:2 form. +const FLAG_PRIVATE: u8 = 0x01; +const FORM_EXTERN_IMPL: u8 = 0x04; + +// `preview$extern_impl` payload flag bits, in the op's declaration order. +const IMPL_ARG_ATTRS: u64 = 1 << 2; +const IMPL_ARCH: u64 = 1 << 6; +const IMPL_FRONTEND_ABI: u64 = 1 << 7; +const IMPL_SYNCHRONIZATION: u64 = 1 << 8; + +/// A type the type table can hold. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +enum Type { + Scalar(ScalarType), + Tile(TileType), + Fragment(FragmentType), + Func(Vec, Vec), +} + +/// Interning tables, filled in first-use order. +struct Tables { + /// The version the module is written at; the fragment type's symbolic + /// shape is encoded differently before 13.6. + version: BytecodeVersion, + strings: Vec, + string_index: HashMap, + types: Vec, + type_index: HashMap, +} + +impl Tables { + fn new(version: BytecodeVersion) -> Self { + Self { + version, + strings: Default::default(), + string_index: Default::default(), + types: Default::default(), + type_index: Default::default(), + } + } + + fn string(&mut self, text: &str) -> u64 { + if let Some(&index) = self.string_index.get(text) { + return index; + } + let index = self.strings.len() as u64; + self.strings.push(text.to_owned()); + self.string_index.insert(text.to_owned(), index); + index + } + + /// Register `ty`, its dependencies first, and return its index. + fn ty(&mut self, ty: &Type) -> u64 { + if let Some(&index) = self.type_index.get(ty) { + return index; + } + match ty { + Type::Scalar(_) => {} + Type::Tile(tile) => { + self.ty(&Type::Scalar(tile.element)); + } + Type::Fragment(fragment) => { + self.ty(&Type::Scalar(fragment.element)); + } + Type::Func(inputs, results) => { + for input in inputs { + self.ty(input); + } + for result in results { + self.ty(result); + } + } + } + let index = self.types.len() as u64; + self.types.push(ty.clone()); + self.type_index.insert(ty.clone(), index); + index + } +} + +/// `varint(count)` then each value as a tagged symbolic integer: the +/// payload of a `SymbolicTuple` and the inline form of a layout's modes. +fn write_int_tuple(w: &mut Writer, values: &[i64]) { + w.varint(values.len() as u64); + for &value in values { + w.varint(ATTR_SYMBOLIC_INT); + w.signed_varint(value); + } +} + +/// A layout attribute: tag, then the shapes and strides tuples inline. +fn write_layout(w: &mut Writer, layout: &Layout) { + w.varint(ATTR_LAYOUT); + write_int_tuple(w, &layout.shapes); + write_int_tuple(w, &layout.strides); +} + +/// The dictionary `{cuda_tile.preview$access = #access<..>}`, tagged. +fn write_access_dictionary(w: &mut Writer, tables: &mut Tables, access: Access) { + w.varint(ATTR_DICTIONARY); + w.varint(1); + w.varint(tables.string(ACCESS_ATTR)); + w.varint(ATTR_ACCESS); + w.varint(access.value()); +} + +/// `arg_attrs`: one access dictionary per argument. Tagged as an array for +/// a declaration (a positional field of the function table), inline (count +/// first, no array tag) inside an implementation's payload. +fn write_arg_attrs(w: &mut Writer, tables: &mut Tables, accesses: &[Access], tagged: bool) { + if tagged { + w.varint(ATTR_ARRAY); + } + w.varint(accesses.len() as u64); + for &access in accesses { + write_access_dictionary(w, tables, access); + } +} + +/// One type-table entry. +fn write_type(w: &mut Writer, tables: &mut Tables, ty: &Type) { + match ty { + Type::Scalar(scalar) => w.varint(u64::from(scalar.tag())), + Type::Tile(tile) => { + w.varint(TAG_TILE); + w.varint(tables.ty(&Type::Scalar(tile.element))); + w.le_var_size_i64(&tile.shape); + } + Type::Func(inputs, results) => { + w.varint(TAG_FUNC); + w.varint(inputs.len() as u64); + for input in inputs { + w.varint(tables.ty(input)); + } + w.varint(results.len() as u64); + for result in results { + w.varint(tables.ty(result)); + } + } + // Tagged symbolic shape, the static shape, element type, memory + // space, both layouts tagged, then the alignment as a tagged symbolic + // integer, always 1. + // + // The symbolic shape's contents depend on the version. The 13.5 + // assembler keeps one integer-valued expression per dimension and + // rejects a rank mismatch ("concrete shape cache rank 1 does not + // match symbolic shape rank 0"). The compiler trunk switched on + // 2026-09-08 to holding only the dynamic dimensions, so a static + // shape writes an empty tuple from 13.6 on. + Type::Fragment(fragment) => { + w.varint(TAG_FRAGMENT); + w.varint(ATTR_SYMBOLIC_TUPLE); + if tables.version < BytecodeVersion::new(13, 6) { + w.varint(fragment.shape.len() as u64); + for &dim in &fragment.shape { + w.varint(ATTR_SYMBOLIC_INT); + w.signed_varint(dim); + } + } else { + w.varint(0); + } + w.le_var_size_i64(&fragment.shape); + w.varint(tables.ty(&Type::Scalar(fragment.element))); + w.varint(fragment.memory_space.value()); + write_layout(w, &fragment.mem_layout); + write_layout(w, &fragment.fragmap); + w.varint(ATTR_SYMBOLIC_INT); + w.signed_varint(1); + } + } +} + +fn declaration_signature(decl: &Declaration) -> Type { + Type::Func( + decl.args + .iter() + .map(|arg| Type::Tile(arg.ty.clone())) + .collect(), + decl.results().map(|ty| Type::Tile(ty.clone())).collect(), + ) +} + +fn implementation_signature(imp: &Implementation) -> Type { + Type::Func( + imp.args + .iter() + .map(|arg| Type::Fragment(arg.fragment.clone())) + .collect(), + Vec::new(), + ) +} + +/// A `preview$func extern` entry: normal form, private, body length zero, +/// then its (empty) constant parameters. +fn write_declaration(w: &mut Writer, tables: &mut Tables, decl: &Declaration, debug: u64) { + w.varint(tables.string(&decl.name)); + w.varint(tables.ty(&declaration_signature(decl))); + w.byte(FLAG_PRIVATE); + w.varint(debug); + let accesses: Vec = decl.args.iter().map(|arg| arg.access).collect(); + write_arg_attrs(w, tables, &accesses, true); + w.varint(0); + w.varint(ATTR_ARRAY); + w.varint(0); +} + +/// A `preview$extern_impl` entry: form bits 01, private, then the op +/// payload in place of a body (flags, symbols, type, the present optional +/// clauses, one empty region) and a body length of zero. +fn write_implementation(w: &mut Writer, tables: &mut Tables, imp: &Implementation, debug: u64) { + w.varint(tables.string(&imp.symbol)); + w.varint(tables.ty(&implementation_signature(imp))); + w.byte(FLAG_PRIVATE | FORM_EXTERN_IMPL); + w.varint(debug); + + let mut flags = IMPL_ARG_ATTRS; + if !imp.arch.is_empty() { + flags |= IMPL_ARCH; + } + if !imp.frontend_abi.is_empty() { + flags |= IMPL_FRONTEND_ABI; + } + if imp.synchronization.is_some() { + flags |= IMPL_SYNCHRONIZATION; + } + w.varint(flags); + w.varint(tables.string(&imp.symbol)); + w.varint(tables.string(&imp.implements)); + w.varint(tables.ty(&implementation_signature(imp))); + let accesses: Vec = imp.args.iter().map(|arg| arg.access).collect(); + write_arg_attrs(w, tables, &accesses, false); + if !imp.arch.is_empty() { + w.varint(imp.arch.len() as u64); + for name in &imp.arch { + w.varint(ATTR_STRING); + w.varint(tables.string(name)); + } + } + if !imp.frontend_abi.is_empty() { + w.varint(imp.frontend_abi.len() as u64); + for (key, value) in &imp.frontend_abi { + w.varint(tables.string(key)); + w.varint(ATTR_STRING); + w.varint(tables.string(value)); + } + } + if let Some(sync) = imp.synchronization { + w.varint(sync.value()); + } + // One region with zero blocks, then the zero body length. + w.varint(1); + w.varint(0); + w.varint(0); +} + +/// Per-function id lists (each holds one id 0, an unknown location) and the +/// single unknown attribute the reader expects a non-empty table to hold. +fn debug_section(functions: usize) -> Vec { + let mut w = Writer::new(); + w.varint(functions as u64); + w.align_to(4); + for index in 0..functions { + w.le_u32(index as u32); + } + w.varint(functions as u64); + w.align_to(8); + for _ in 0..functions { + w.le_u64(0); + } + w.varint(1); + w.align_to(4); + w.le_u32(0); + w.byte(0); + w.into_bytes() +} + +fn type_section(tables: &mut Tables) -> Vec { + let mut w = Writer::new(); + let count = tables.types.len(); + let offsets = begin_offset_table(&mut w, count); + let start = w.len(); + // Writing an entry can only reference types registered before it, so + // the table does not grow while it is serialized. + let types = tables.types.clone(); + for (index, ty) in types.iter().enumerate() { + w.patch_u32(offsets + 4 * index, (w.len() - start) as u32); + write_type(&mut w, tables, ty); + } + debug_assert_eq!(tables.types.len(), count); + w.into_bytes() +} + +fn string_section(tables: &Tables) -> Vec { + let mut w = Writer::new(); + let offsets = begin_offset_table(&mut w, tables.strings.len()); + let mut running = 0u32; + for (index, text) in tables.strings.iter().enumerate() { + w.patch_u32(offsets + 4 * index, running); + w.bytes(text.as_bytes()); + running += text.len() as u32; + } + w.into_bytes() +} + +fn write_module( + version: BytecodeVersion, + entries: usize, + mut write_entries: impl FnMut(&mut Writer, &mut Tables), +) -> Result, Error> { + if version < BytecodeVersion::PREVIEW_FIRST { + return Err(Error::Version(version)); + } + let mut tables = Tables::new(version); + let mut functions = Writer::new(); + functions.varint(entries as u64); + write_entries(&mut functions, &mut tables); + functions.align_to(8); + + let mut out = MAGIC.to_vec(); + out.push(version.major); + out.push(version.minor); + out.extend_from_slice(&CONTAINS_PREVIEW.to_le_bytes()); + write_section(&mut out, Section::Func, &functions.into_bytes(), 8); + write_section(&mut out, Section::Debug, &debug_section(entries), 8); + write_section(&mut out, Section::Type, &type_section(&mut tables), 4); + write_section(&mut out, Section::String, &string_section(&tables), 4); + out.push(Section::EndOfBytecode as u8); + Ok(out) +} + +impl Library { + /// The interface half: every declaration as a `preview$func extern`. + /// + /// `version` must be the one the consuming assembler calls current; + /// preview content is readable at no other version. Anything below + /// 13.5 is an error. + pub fn interface_bytecode(&self, version: BytecodeVersion) -> Result, Error> { + write_module(version, self.declarations.len(), |w, tables| { + for (index, decl) in self.declarations.iter().enumerate() { + write_declaration(w, tables, decl, index as u64 + 1); + } + }) + } + + /// The implementation half: every implementation as a + /// `preview$extern_impl`. + pub fn implementation_bytecode(&self, version: BytecodeVersion) -> Result, Error> { + write_module(version, self.implementations.len(), |w, tables| { + for (index, imp) in self.implementations.iter().enumerate() { + write_implementation(w, tables, imp, index as u64 + 1); + } + }) + } +} diff --git a/cuda-oxide/crates/ptx-parse/README.md b/cuda-oxide/crates/ptx-parse/README.md index 91721b81c0..7a18cfb525 100644 --- a/cuda-oxide/crates/ptx-parse/README.md +++ b/cuda-oxide/crates/ptx-parse/README.md @@ -89,4 +89,4 @@ re-parsing. ## License -Apache-2.0. See [LICENSE](https://github.com/NVlabs/cuda-oxide/blob/main/LICENSE). +Apache-2.0. See [LICENSE](https://github.com/NVIDIA/cuda-rust/blob/main/LICENSE). diff --git a/cuda-oxide/crates/ptx-schedule/README.md b/cuda-oxide/crates/ptx-schedule/README.md index 872b8e858b..14caec6ba6 100644 --- a/cuda-oxide/crates/ptx-schedule/README.md +++ b/cuda-oxide/crates/ptx-schedule/README.md @@ -99,6 +99,14 @@ ptx-schedule kernel.ptx --seed 7 --intensity 0.5 -o perturbed.ptx \ --decisions-json decisions.json ``` +`--decisions-json` must name a file other than the input and `-o`, and it +cannot be a directory. The binary resolves all three paths before it writes +anything: it follows symbolic links, including one whose target does not exist +yet, and on Unix it compares existing files by device and inode (the identity +on disk that a hard link shares with its original). A collision is refused +with `refusing to overwrite PTX` and nothing is written. Rewriting the PTX in +place, with `-o` equal to the input and a separate report file, stays allowed. + ## License -Apache-2.0. See [LICENSE](https://github.com/NVlabs/cuda-oxide/blob/main/LICENSE). +Apache-2.0. See [LICENSE](https://github.com/NVIDIA/cuda-rust/blob/main/LICENSE). diff --git a/cuda-oxide/crates/ptx-schedule/src/main.rs b/cuda-oxide/crates/ptx-schedule/src/main.rs index ebbff3ad8b..9c9c7efaec 100644 --- a/cuda-oxide/crates/ptx-schedule/src/main.rs +++ b/cuda-oxide/crates/ptx-schedule/src/main.rs @@ -6,7 +6,7 @@ use clap::{ArgGroup, Parser}; use ptx_schedule::{InjectionOptions, analyze_ptx, perturb_ptx}; use std::fs; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; #[derive(Debug, Parser)] #[command( @@ -64,7 +64,8 @@ struct Cli { )] focus: Option, - /// Save the injection decisions as JSON. + /// Save the injection decisions as JSON. Must differ from the input + /// and --output; a collision is refused before anything is written. #[arg( long, value_name = "FILE", @@ -74,6 +75,109 @@ struct Cli { decisions_json: Option, } +/// Longest chain of symbolic links followed for one destination. Linux +/// follows at most 40 links in one path lookup, so a chain the kernel accepts +/// is never refused here; the loop below runs one extra time to resolve the +/// final non-link component. +const MAX_SYMLINK_FOLLOWS: usize = 40; + +/// Canonical form of a path that may not exist yet, so that two names for one +/// destination compare equal. An existing path is canonicalized. A dangling +/// symbolic link (a link whose target does not exist yet) is followed to its +/// target, because a write through the link lands there. Otherwise the +/// existing parent is canonicalized and the file name is appended, so `./` +/// and `../` forms compare equal. `label` names the flag in error messages. +/// Hard links keep distinct canonical paths; `same_file` covers them. +fn canonical_destination(path: &Path, label: &str) -> Result> { + let mut absolute = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir()?.join(path) + }; + for _ in 0..=MAX_SYMLINK_FOLLOWS { + if let Ok(canonical) = fs::canonicalize(&absolute) { + return Ok(canonical); + } + let (Some(file_name), Some(parent)) = (absolute.file_name(), absolute.parent()) else { + return Err(format!("{label} path has no file name: {}", absolute.display()).into()); + }; + if !absolute.is_symlink() { + let parent = fs::canonicalize(parent).map_err(|err| { + format!("cannot resolve {label} path {}: {err}", absolute.display()) + })?; + return Ok(parent.join(file_name)); + } + let target = fs::read_link(&absolute) + .map_err(|err| format!("cannot resolve {label} path {}: {err}", absolute.display()))?; + absolute = if target.is_absolute() { + target + } else { + parent.join(target) + }; + } + Err(format!( + "{label} path has too many levels of symbolic links: {}", + path.display() + ) + .into()) +} + +/// Whether two existing paths name one file on disk. A hard link has its own +/// canonical path but shares the device and inode of the file it aliases, so +/// the path comparison alone misses it. A path that does not exist yet +/// answers false, and the path comparison still applies to it. +#[cfg(unix)] +fn same_file(a: &Path, b: &Path) -> bool { + use std::os::unix::fs::MetadataExt; + match (fs::metadata(a), fs::metadata(b)) { + (Ok(a), Ok(b)) => a.dev() == b.dev() && a.ino() == b.ino(), + _ => false, + } +} + +/// Only Unix exposes a file identity through `std`; other platforms keep the +/// canonical path comparison. +#[cfg(not(unix))] +fn same_file(_a: &Path, _b: &Path) -> bool { + false +} + +/// Refuses a `--decisions-json` destination that would replace the PTX input +/// or `--output`, or that is a directory, before anything is written. +fn reject_decisions_clobber( + input: &Path, + output: &Path, + decisions: &Path, +) -> Result<(), Box> { + let decisions_path = canonical_destination(decisions, "--decisions-json")?; + if decisions_path.is_dir() { + return Err(format!( + "--decisions-json names a directory: {}", + decisions_path.display() + ) + .into()); + } + let input_path = canonical_destination(input, "the PTX input")?; + let output_path = canonical_destination(output, "--output")?; + let same_as_input = decisions_path == input_path || same_file(&decisions_path, &input_path); + let same_as_output = decisions_path == output_path || same_file(&decisions_path, &output_path); + if same_as_input || same_as_output { + let which = if same_as_input && same_as_output { + "the PTX input and --output" + } else if same_as_input { + "the PTX input" + } else { + "--output" + }; + return Err(format!( + "--decisions-json resolves to {}, which is {which}; refusing to overwrite PTX", + decisions_path.display() + ) + .into()); + } + Ok(()) +} + fn main() -> Result<(), Box> { let cli = Cli::parse(); let source = fs::read_to_string(&cli.input)?; @@ -100,9 +204,12 @@ fn main() -> Result<(), Box> { let output = cli .output .expect("clap requires either --list-sites or --output"); + if let Some(decisions) = &cli.decisions_json { + reject_decisions_clobber(&cli.input, &output, decisions)?; + } let rewrite = perturb_ptx(&source, &options)?; - fs::write(output, &rewrite.ptx)?; - if let Some(path) = cli.decisions_json { + fs::write(&output, &rewrite.ptx)?; + if let Some(path) = &cli.decisions_json { fs::write(path, serde_json::to_string_pretty(&rewrite.report)?)?; } println!( diff --git a/cuda-oxide/crates/ptx-schedule/tests/cli.rs b/cuda-oxide/crates/ptx-schedule/tests/cli.rs index 5887f3c4e4..ec29b40578 100644 --- a/cuda-oxide/crates/ptx-schedule/tests/cli.rs +++ b/cuda-oxide/crates/ptx-schedule/tests/cli.rs @@ -35,12 +35,21 @@ impl Fixture { } fn run(&self, args: &[&str]) -> Output { + self.run_in(&self.dir, args) + } + + fn run_in(&self, cwd: &Path, args: &[&str]) -> Output { Command::new(env!("CARGO_BIN_EXE_ptx-schedule")) + .current_dir(cwd) .arg(&self.input) .args(args) .output() .unwrap() } + + fn bytes(&self) -> Vec { + fs::read(&self.input).unwrap() + } } impl Drop for Fixture { @@ -75,6 +84,11 @@ fn help_shows_the_two_exclusive_forms() { let stdout = String::from_utf8_lossy(&output.stdout); assert!(stdout.contains("ptx-schedule --list-sites")); assert!(stdout.contains("ptx-schedule -o [OPTIONS]")); + let flattened = stdout.split_whitespace().collect::>().join(" "); + assert!( + flattened.contains("Must differ from the input and --output"), + "{stdout}" + ); } #[test] @@ -107,3 +121,273 @@ fn both_valid_modes_execute() { assert!(rewritten.status.success(), "{rewritten:?}"); assert!(output_path.is_file()); } + +fn assert_rejected(output: &Output, fragment: &str) { + assert!( + !output.status.success(), + "run should be rejected, got {output:?}" + ); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + !stdout.contains("ptx-schedule: seed="), + "success summary must not claim a write that was rejected: {stdout}" + ); + assert!(stderr(output).contains(fragment), "{}", stderr(output)); +} + +/// A `--decisions-json` destination that is the PTX input or `--output`. +fn assert_decisions_rejected(output: &Output) { + assert_rejected(output, "refusing to overwrite PTX"); +} + +#[test] +fn decisions_json_exact_output_collision_is_rejected() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let output_path = fixture.output(); + fs::write(&output_path, "ORIGINAL-PTX").unwrap(); + let path = output_path.to_str().unwrap(); + let output = fixture.run(&["--seed", "7", "--output", path, "--decisions-json", path]); + assert_decisions_rejected(&output); + assert_eq!(fixture.bytes(), before); + assert_eq!(fs::read_to_string(&output_path).unwrap(), "ORIGINAL-PTX"); +} + +#[test] +fn decisions_json_dot_alias_of_output_is_rejected() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let output_path = fixture.output(); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + "output.ptx", + "--decisions-json", + "./output.ptx", + ]); + assert_decisions_rejected(&output); + assert_eq!(fixture.bytes(), before); + assert!( + !output_path.exists(), + "rejected run must not create PTX output" + ); +} + +#[test] +fn decisions_json_dotdot_alias_of_output_is_rejected() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let nested = fixture.dir.join("nested"); + fs::create_dir_all(&nested).unwrap(); + let output_path = fixture.output(); + let output = fixture.run_in( + &nested, + &[ + "--seed", + "7", + "--output", + "../output.ptx", + "--decisions-json", + "../nested/../output.ptx", + ], + ); + assert_decisions_rejected(&output); + assert_eq!(fixture.bytes(), before); + assert!( + !output_path.exists(), + "rejected run must not create PTX output" + ); +} + +#[test] +fn decisions_json_equal_to_input_leaves_input_unchanged() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let output_path = fixture.output(); + let input = fixture.input.to_str().unwrap(); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + output_path.to_str().unwrap(), + "--decisions-json", + input, + ]); + assert_decisions_rejected(&output); + assert_eq!(fixture.bytes(), before); + assert!( + !output_path.exists(), + "rejected run must not write --output" + ); + + let dotted = fixture.run(&[ + "--seed", + "7", + "--output", + "output.ptx", + "--decisions-json", + "./input.ptx", + ]); + assert_decisions_rejected(&dotted); + assert_eq!(fixture.bytes(), before); + assert!( + !output_path.exists(), + "rejected run must not write --output" + ); +} + +#[test] +fn inplace_ptx_rewrite_allows_a_distinct_decisions_json() { + let fixture = Fixture::new(); + let decisions = fixture.dir.join("decisions.json"); + let input = fixture.input.to_str().unwrap(); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + input, + "--decisions-json", + decisions.to_str().unwrap(), + ]); + assert!(output.status.success(), "{output:?}"); + assert!(String::from_utf8_lossy(&output.stdout).contains("ptx-schedule: seed=")); + let ptx = fs::read_to_string(&fixture.input).unwrap(); + assert!(ptx.contains("nanosleep.u32"), "{ptx}"); + assert!(!ptx.trim_start().starts_with('{'), "{ptx}"); + let json = fs::read_to_string(&decisions).unwrap(); + assert!(json.contains("\"sites_injected\""), "{json}"); +} + +#[cfg(unix)] +#[test] +fn decisions_json_dangling_symlink_to_output_is_rejected() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let output_path = fixture.output(); + // The link's target does not exist yet; a write through the link creates it. + std::os::unix::fs::symlink("output.ptx", fixture.dir.join("link.json")).unwrap(); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + "output.ptx", + "--decisions-json", + "link.json", + ]); + assert_rejected(&output, "which is --output; refusing to overwrite PTX"); + assert_eq!(fixture.bytes(), before); + assert!( + !output_path.exists(), + "rejected run must not create PTX output" + ); +} + +#[cfg(unix)] +#[test] +fn output_dangling_symlink_to_decisions_target_is_rejected() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let target = fixture.dir.join("target.ptx"); + std::os::unix::fs::symlink("target.ptx", fixture.dir.join("outlink.ptx")).unwrap(); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + "outlink.ptx", + "--decisions-json", + "target.ptx", + ]); + assert_rejected(&output, "which is --output; refusing to overwrite PTX"); + assert_eq!(fixture.bytes(), before); + assert!( + !target.exists(), + "rejected run must not create the link target" + ); +} + +#[cfg(unix)] +#[test] +fn decisions_json_hard_link_of_input_is_rejected() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let output_path = fixture.output(); + fs::hard_link(&fixture.input, fixture.dir.join("alias.json")).unwrap(); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + "output.ptx", + "--decisions-json", + "alias.json", + ]); + assert_rejected(&output, "which is the PTX input; refusing to overwrite PTX"); + assert_eq!(fixture.bytes(), before); + assert!( + !output_path.exists(), + "rejected run must not create PTX output" + ); +} + +#[cfg(unix)] +#[test] +fn decisions_json_hard_link_of_output_is_rejected() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let output_path = fixture.output(); + fs::write(&output_path, "ORIGINAL-PTX").unwrap(); + fs::hard_link(&output_path, fixture.dir.join("alias.json")).unwrap(); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + "output.ptx", + "--decisions-json", + "alias.json", + ]); + assert_rejected(&output, "which is --output; refusing to overwrite PTX"); + assert_eq!(fixture.bytes(), before); + assert_eq!(fs::read_to_string(&output_path).unwrap(), "ORIGINAL-PTX"); +} + +#[test] +fn decisions_json_directory_is_rejected_before_the_ptx_write() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let output_path = fixture.output(); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + "output.ptx", + "--decisions-json", + ".", + ]); + assert_rejected(&output, "--decisions-json names a directory"); + assert_eq!(fixture.bytes(), before); + assert!( + !output_path.exists(), + "rejected run must not create PTX output" + ); +} + +#[test] +fn output_without_a_file_name_is_blamed_on_output() { + let fixture = Fixture::new(); + let before = fixture.bytes(); + let bad_output = fixture.dir.join("missing").join(".."); + let output = fixture.run(&[ + "--seed", + "7", + "--output", + bad_output.to_str().unwrap(), + "--decisions-json", + "decisions.json", + ]); + assert_rejected(&output, "--output path has no file name"); + assert_eq!(fixture.bytes(), before); + assert!( + !fixture.dir.join("decisions.json").exists(), + "rejected run must not write the report" + ); +} diff --git a/cuda-oxide/crates/reserved-oxide-symbols/README.md b/cuda-oxide/crates/reserved-oxide-symbols/README.md index cc8e0ad98d..0acc711a93 100644 --- a/cuda-oxide/crates/reserved-oxide-symbols/README.md +++ b/cuda-oxide/crates/reserved-oxide-symbols/README.md @@ -17,7 +17,7 @@ Every symbol prefix here contains the magic component `246e25db_`, which is exists purely to make accidental collisions impossible — a user is never going to write `fn cuda_oxide_codegen_v1_cuda_oxide_kernel_246e25db_foo()` by accident. -The crate owns these ten symbol constants. The `naming-guard` CI job keeps +The crate owns these eleven symbol constants. The `naming-guard` CI job keeps reserved-prefix literals out of their consumers: | Constant | Value | What it names | @@ -27,6 +27,7 @@ reserved-prefix literals out of their consumers: | `DEVICE_EXTERN_PREFIX` | `cuda_oxide_device_extern_246e25db_` | `extern` device declarations | | `INSTANTIATE_PREFIX` | `cuda_oxide_instantiate_246e25db_` | forced generic instantiations | | `CONSTANT_PREFIX` | `cuda_oxide_const_246e25db_` | `#[constant]` statics, for `cuModuleGetGlobal` lookup | +| `TILE_EXPORT_META_PREFIX` | `cuda_oxide_tile_export_246e25db_` | host-only `#[tile_export]` metadata functions whose body returns the export's JSON description | | `KERNEL_SCOPE_LOCAL` | `cuda_oxide_kernel_scope_246e25db` | the hidden thread-index scope token injected by `#[kernel]` / `#[device]` | | `ARTIFACT_ANCHOR_PREFIX` | `cuda_oxide_artifact_anchor_246e25db_` | the link anchor pinning a crate's `.oxart` section into the binary | | `PTX_MERGE_REQUIRED_PREFIX` | `cuda_oxide_ptx_merge_required_246e25db_` | statics marking modules whose generic kernels need PTX-bundle merging | diff --git a/cuda-oxide/crates/reserved-oxide-symbols/src/lib.rs b/cuda-oxide/crates/reserved-oxide-symbols/src/lib.rs index f21877ee87..010bc715d2 100644 --- a/cuda-oxide/crates/reserved-oxide-symbols/src/lib.rs +++ b/cuda-oxide/crates/reserved-oxide-symbols/src/lib.rs @@ -161,6 +161,22 @@ pub const INSTANTIATE_PREFIX: &str = "cuda_oxide_instantiate_246e25db_"; /// `set_coeffs` methods generated by `#[cuda_module]` resolve that exact name. pub const CONSTANT_PREFIX: &str = "cuda_oxide_const_246e25db_"; +/// Prefix of the host-only metadata function emitted by `#[tile_export]`. +/// +/// `tile_export_meta_symbol("relu_impl")` produces +/// `cuda_oxide_tile_export_246e25db_relu_impl`. The macro emits one such +/// `#[no_mangle] fn` per exported function; its body returns the JSON +/// description of the export as a `&'static str`. The codegen backend reads +/// that constant out of the function's MIR, writes the collected descriptions +/// to `.tileexports.json` next to the device artifact, and never +/// compiles the function itself for the device. +/// +/// The prefix contains neither [`LEGACY_KERNEL_PREFIX`] nor +/// [`LEGACY_DEVICE_PREFIX`] as a substring, so [`is_kernel_symbol`] and +/// [`is_device_symbol`] can never classify a metadata function as device +/// code. A unit test below pins that property. +pub const TILE_EXPORT_META_PREFIX: &str = "cuda_oxide_tile_export_246e25db_"; + /// Local binding name injected by `#[kernel]` and `#[device]` for the /// hidden thread-index scope token. pub const KERNEL_SCOPE_LOCAL: &str = "cuda_oxide_kernel_scope_246e25db"; @@ -266,6 +282,20 @@ pub fn constant_symbol(base: &str) -> String { format!("{CONSTANT_PREFIX}{base}") } +/// Build the `#[tile_export]` metadata-function symbol for an implementation +/// symbol. +/// +/// ``` +/// use reserved_oxide_symbols::tile_export_meta_symbol; +/// assert_eq!( +/// tile_export_meta_symbol("relu_impl"), +/// "cuda_oxide_tile_export_246e25db_relu_impl", +/// ); +/// ``` +pub fn tile_export_meta_symbol(base: &str) -> String { + format!("{TILE_EXPORT_META_PREFIX}{base}") +} + /// Build the compiler marker associated with one generic kernel. pub fn ptx_merge_required_marker(base: &str) -> String { let mut symbol = String::from(PTX_MERGE_REQUIRED_PREFIX); @@ -459,6 +489,26 @@ pub fn is_constant_symbol(name: &str) -> bool { name.contains(CONSTANT_PREFIX) } +/// Returns `true` if the last path component of `name` is a `#[tile_export]` +/// metadata function. +/// +/// Matches the final `::` component only, like +/// [`is_ptx_merge_required_marker`]: the backend looks these functions up by +/// definition path, and an enclosing module named after the prefix must not +/// turn its contents into metadata functions. +/// +/// ``` +/// use reserved_oxide_symbols::{is_tile_export_meta_symbol, tile_export_meta_symbol}; +/// assert!(is_tile_export_meta_symbol(&tile_export_meta_symbol("relu_impl"))); +/// assert!(is_tile_export_meta_symbol( +/// "tilelib::cuda_oxide_tile_export_246e25db_relu_impl" +/// )); +/// assert!(!is_tile_export_meta_symbol("relu_impl")); +/// ``` +pub fn is_tile_export_meta_symbol(name: &str) -> bool { + tile_export_meta_base_name(name).is_some() +} + /// Strip the kernel prefix from a possibly-FQDN symbol name. /// /// Returns the part of `name` after [`KERNEL_PREFIX`], or `None` if `name` @@ -551,6 +601,8 @@ pub fn display_name(name: &str) -> Option { Some(format!("{base} (device)")) } else if let Some(base) = constant_base_name(name) { Some(format!("{base} (constant)")) + } else if let Some(base) = tile_export_meta_base_name(name) { + Some(format!("{base} (tile export metadata)")) } else { instantiate_base_name(name).map(|base| format!("{base} (instantiate helper)")) } @@ -589,6 +641,29 @@ pub fn constant_base_name(name: &str) -> Option<&str> { .map(|pos| &name[pos + CONSTANT_PREFIX.len()..]) } +/// Strip the `#[tile_export]` metadata prefix from the last path component +/// of `name`, returning the implementation symbol it describes. +/// +/// ``` +/// use reserved_oxide_symbols::tile_export_meta_base_name; +/// assert_eq!( +/// tile_export_meta_base_name("cuda_oxide_tile_export_246e25db_relu_impl"), +/// Some("relu_impl"), +/// ); +/// assert_eq!( +/// tile_export_meta_base_name("tilelib::kernels::cuda_oxide_tile_export_246e25db_scale_impl"), +/// Some("scale_impl"), +/// ); +/// assert_eq!(tile_export_meta_base_name("relu_impl"), None); +/// assert_eq!(tile_export_meta_base_name("cuda_oxide_tile_export_246e25db_"), None); +/// ``` +pub fn tile_export_meta_base_name(name: &str) -> Option<&str> { + name.rsplit("::") + .next() + .and_then(|component| component.strip_prefix(TILE_EXPORT_META_PREFIX)) + .filter(|base| !base.is_empty()) +} + // ============================================================================ // Tests — pin the constant values, verify mutual-exclusion, round-trip // ============================================================================ @@ -617,6 +692,7 @@ mod tests { assert_eq!(DEVICE_EXTERN_PREFIX, "cuda_oxide_device_extern_246e25db_"); assert_eq!(INSTANTIATE_PREFIX, "cuda_oxide_instantiate_246e25db_"); assert_eq!(CONSTANT_PREFIX, "cuda_oxide_const_246e25db_"); + assert_eq!(TILE_EXPORT_META_PREFIX, "cuda_oxide_tile_export_246e25db_"); assert_eq!( PTX_MERGE_REQUIRED_PREFIX, "cuda_oxide_ptx_merge_required_246e25db_" @@ -660,6 +736,7 @@ mod tests { DEVICE_EXTERN_PREFIX, INSTANTIATE_PREFIX, CONSTANT_PREFIX, + TILE_EXPORT_META_PREFIX, PTX_MERGE_REQUIRED_PREFIX, ARTIFACT_ANCHOR_PREFIX, KERNEL_SCOPE_LOCAL, @@ -683,6 +760,45 @@ mod tests { assert!(!DEVICE_PREFIX.contains(DEVICE_EXTERN_PREFIX)); } + /// The backend finds `#[tile_export]` metadata functions by definition + /// path and must never hand one to device codegen. That holds because + /// the prefix contains no kernel or device prefix, current or legacy, + /// so the device predicates cannot match a metadata symbol, and because + /// the metadata predicate only looks at the final path component. + #[test] + fn tile_export_meta_symbols_are_neither_kernels_nor_device_functions() { + for device_like in [ + KERNEL_PREFIX, + LEGACY_KERNEL_PREFIX, + DEVICE_PREFIX, + LEGACY_DEVICE_PREFIX, + DEVICE_EXTERN_PREFIX, + ] { + assert!(!TILE_EXPORT_META_PREFIX.contains(device_like)); + assert!(!device_like.contains(TILE_EXPORT_META_PREFIX)); + } + + let meta = tile_export_meta_symbol("relu_impl"); + assert!(is_tile_export_meta_symbol(&meta)); + assert!(!is_kernel_symbol(&meta)); + assert!(!is_device_symbol(&meta)); + assert!(!is_device_extern_symbol(&meta)); + assert!(!is_constant_symbol(&meta)); + assert_eq!(tile_export_meta_base_name(&meta), Some("relu_impl")); + + // A device function and a kernel are not metadata functions. + assert!(!is_tile_export_meta_symbol(&device_symbol("relu_impl"))); + assert!(!is_tile_export_meta_symbol(&kernel_symbol("relu"))); + + // Only the final path component counts. + assert!(is_tile_export_meta_symbol(&format!( + "tilelib::kernels::{meta}" + ))); + assert!(!is_tile_export_meta_symbol(&format!("{meta}::child"))); + assert!(!is_tile_export_meta_symbol(&format!("prefix_{meta}"))); + assert!(!is_tile_export_meta_symbol(TILE_EXPORT_META_PREFIX)); + } + #[test] fn ptx_merge_markers_are_per_kernel_and_sanitized() { let marker = ptx_merge_required_marker("nested::map-value"); @@ -738,6 +854,10 @@ mod tests { ); assert_eq!(instantiate_base_name(&instantiate_symbol(base)), Some(base),); assert_eq!(constant_base_name(&constant_symbol(base)), Some(base)); + assert_eq!( + tile_export_meta_base_name(&tile_export_meta_symbol(base)), + Some(base) + ); } } @@ -851,6 +971,10 @@ mod tests { display_name(&constant_symbol("COEFFS")).as_deref(), Some("COEFFS (constant)"), ); + assert_eq!( + display_name(&tile_export_meta_symbol("relu_impl")).as_deref(), + Some("relu_impl (tile export metadata)"), + ); assert_eq!(display_name("std::vec::Vec::new"), None); } @@ -865,12 +989,17 @@ mod tests { "cuda_oxide_device_extern_evil", "cuda_oxide_instantiate_evil", "cuda_oxide_const_evil", + "cuda_oxide_tile_export_evil", ] { assert!(!is_kernel_symbol(evil), "unexpected match: {evil}"); assert!(!is_device_symbol(evil), "unexpected match: {evil}"); assert!(!is_device_extern_symbol(evil), "unexpected match: {evil}"); assert!(!is_instantiate_symbol(evil), "unexpected match: {evil}"); assert!(!is_constant_symbol(evil), "unexpected match: {evil}"); + assert!( + !is_tile_export_meta_symbol(evil), + "unexpected match: {evil}" + ); assert_eq!(display_name(evil), None); } } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/.gitignore b/cuda-oxide/crates/rustc-codegen-cuda/.gitignore index f528657c89..e80f489098 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/.gitignore +++ b/cuda-oxide/crates/rustc-codegen-cuda/.gitignore @@ -7,6 +7,14 @@ rustc-ice-*.txt # Generated PTX/LLVM files *.ptx *.ll +# Tile export descriptions written next to them +*.tileexports.json +# Tile library descriptions and the library itself (cargo oxide tilelib) +*.tilebc +*.tilelib +*.tilelib.mlir +# Mock Tile-compiler callers that tile library examples generate when they run +examples/*/mock/ # Implementation log (internal notes) IMPLEMENTATION.md diff --git a/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock index 0c08699d3f..512dd99178 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/Cargo.lock @@ -31,9 +31,9 @@ checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" [[package]] name = "awint" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d48b7360a36d335663e8f20b7f439029debf52b5a0a749d6e936405f25045288" +checksum = "b2a8f298efd080b0d3e1ee1723b5d88ecc0d78805fad347e09b76dbd4dadbe37" dependencies = [ "awint_core", "awint_dag", @@ -44,9 +44,9 @@ dependencies = [ [[package]] name = "awint_core" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8280079e78217ace721501f35dbf551dadf0ab63863504169316cea1bbac872" +checksum = "c24d226d9faa090abec192c8ad1ea74635e6cc94000c72a95ceb60cf691428c3" dependencies = [ "awint_internals", "const_fn", @@ -54,9 +54,9 @@ dependencies = [ [[package]] name = "awint_dag" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73977ba1786a5e127272f08b865b60f5d548576a0925de1daa5703cc9ba78a8d" +checksum = "dc5377016f3077b18304da4233a28626ea069047ca84c9bc8c8b81b659f11032" dependencies = [ "awint_ext", "awint_macro_internals", @@ -66,9 +66,9 @@ dependencies = [ [[package]] name = "awint_ext" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef1b30d3640d9f94dd9f55f655773c5238c11be6f95ab36d24308b7ae34c9bbc" +checksum = "01ff843cdefdd48fbe9bc11292a902b7921ff37433ed4124fa02136f83186f58" dependencies = [ "awint_core", "const_fn", @@ -76,18 +76,18 @@ dependencies = [ [[package]] name = "awint_internals" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "173937e3f4e233d93362fc1e28ab23808cb31671a1923033706cb44ed4e5d10c" +checksum = "6209d7873685ac95ee41c4b63fe29bff9b2e1106de65d05b106b2b769c178575" dependencies = [ "const_fn", ] [[package]] name = "awint_macro_internals" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2c6a23f64f14c3b566c2c04124015f78c5fb93e8275a6574e0118747ee60de6" +checksum = "6b519a9cd37bf73e44e29b3779ca90246412b56db4af4e96d959f5082bb8bdcc" dependencies = [ "awint_ext", "proc-macro2", @@ -96,19 +96,18 @@ dependencies = [ [[package]] name = "awint_macros" -version = "0.18.1" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16e50a2862c6002a424fc9bb4982ec33ab679cb4ad04a174e2a52c2ecc3ea4c7" +checksum = "518a133c2163e8f021f708d7278af27dd6c1bcee96ea672a2d825b13a6931818" dependencies = [ - "awint_internals", "awint_macro_internals", ] [[package]] name = "bitflags" -version = "2.10.0" +version = "2.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" [[package]] name = "block-buffer" @@ -139,9 +138,9 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "combine" -version = "4.6.7" +version = "4.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" dependencies = [ "bytes", "memchr", @@ -155,9 +154,9 @@ checksum = "2f8a2ca5ac02d09563609681103aada9e1777d54fc57a5acd7a41404f9c93b6e" [[package]] name = "convert_case" -version = "0.11.0" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "affbf0190ed2caf063e3def54ff444b449371d55c58e513a95ab98eca50adb49" +checksum = "1af709f1f33454bf52eadfc8c78b3b9ef9cb26fb54d16dc9cd9a7299f899fd1b" dependencies = [ "unicode-segmentation", ] @@ -205,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" @@ -351,9 +393,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" dependencies = [ "allocator-api2", "equivalent", @@ -362,9 +404,9 @@ dependencies = [ [[package]] name = "hi_sparse_bitset" -version = "0.8.0" +version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad02fc3cac722c3c44c3ed708f6653e7d149b7912f7342fc1e8530455f84f0eb" +checksum = "682097abdf22a902b7b2b5f2824e773691507289873c3ea0abbc5d4c85f05ec0" dependencies = [ "wide", ] @@ -380,19 +422,19 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.14.0" +version = "2.14.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" dependencies = [ "equivalent", - "hashbrown 0.17.0", + "hashbrown 0.17.1", ] [[package]] name = "inventory" -version = "0.3.24" +version = "0.3.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4f0c30c76f2f4ccee3fe55a2435f691ca00c0e4bd87abe4f4a851b1d4dac39b" +checksum = "6928282826c822ad91bf1c9a1cb90a30ba1c26770749929b4656cd6be829cd7c" dependencies = [ "rustversion", ] @@ -430,22 +472,22 @@ dependencies = [ [[package]] name = "linkme" -version = "0.3.35" +version = "0.3.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e3283ed2d0e50c06dd8602e0ab319bb048b6325d0bba739db64ed8205179898" +checksum = "3045e122bd98aef8ec3ad58ce84f0791f64e70163d1a02710af4aa11a4d54cc5" dependencies = [ "linkme-impl", ] [[package]] name = "linkme-impl" -version = "0.3.35" +version = "0.3.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5cec0ec4228b4853bb129c84dbf093a27e6c7a20526da046defc334a1b017f7" +checksum = "77060ebe535362c3da75682cd17b0431017b6e7c5661e714fc69a7ad017d1301" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -482,9 +524,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.29" +version = "0.4.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" [[package]] name = "memchr" @@ -498,6 +540,8 @@ version = "0.2.1" dependencies = [ "combine", "cuda-oxide-codegen", + "cute-layout", + "dialect-cute", "dialect-iket", "dialect-mir", "dialect-nvvm", @@ -576,9 +620,9 @@ checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" [[package]] name = "once_vec" -version = "0.3.0" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a6298074a1742b584aaca148a6c61ecb27f0b55fdf237c235939c84924377a1" +checksum = "0b61fbf7ec37342ce171d9f5e3b82be74bda6c4e84e3b28338a03d553913e6aa" [[package]] name = "oxide-artifacts" @@ -589,16 +633,26 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "pliron" -version = "0.18.0" -source = "git+https://github.com/pliron-org/pliron?rev=06d99a3df938394583715aacf5eadb65b4791c09#06d99a3df938394583715aacf5eadb65b4791c09" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7395b93d3fc99ee1923027fedead143eb63e4c76b6f6f6d79ebfac6698956317" dependencies = [ "awint", "combine", "downcast-rs", "dyn-clone", - "hashbrown 0.17.0", + "hashbrown 0.17.1", "hi_sparse_bitset", "indexmap", "inventory", @@ -618,8 +672,9 @@ dependencies = [ [[package]] name = "pliron-derive" -version = "0.18.0" -source = "git+https://github.com/pliron-org/pliron?rev=06d99a3df938394583715aacf5eadb65b4791c09#06d99a3df938394583715aacf5eadb65b4791c09" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "245f8ea25d8ddcaf04e3342be709b4f257cef2c2c15b99d263955d6f073796cd" dependencies = [ "combine", "convert_case", @@ -627,13 +682,14 @@ dependencies = [ "proc-macro2", "quote", "rustc-hash", - "syn 3.0.3", + "syn", ] [[package]] name = "pliron-llvm" -version = "0.18.0" -source = "git+https://github.com/pliron-org/pliron?rev=06d99a3df938394583715aacf5eadb65b4791c09#06d99a3df938394583715aacf5eadb65b4791c09" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b027d3174757ccc3a96ae971aa4e7904439a3082a3550bc01361369c3ab073f" dependencies = [ "bitflags", "log", @@ -641,11 +697,22 @@ 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.106" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] @@ -656,9 +723,9 @@ version = "0.1.0" [[package]] name = "quote" -version = "1.0.44" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] @@ -671,15 +738,15 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "recasting" -version = "0.2.1" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70aeee1a07922bfe7c79e148553dee262248d2493d9778eee7551b9b5d7cc651" +checksum = "dc3f51d3a1399f13fa081d97345d0fdf08597fe78558fc8539d0e67412f83b26" [[package]] name = "regex" -version = "1.12.2" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843bc0191f75f3e22651ae5f1e72939ab2f72a4bc30fa80a066bd66edefc24d4" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" dependencies = [ "aho-corasick", "memchr", @@ -689,9 +756,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.13" +version = "0.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5276caf25ac86c8d810222b3dbb938e512c55c6831a10f3e6ed1c93b84041f1c" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" dependencies = [ "aho-corasick", "memchr", @@ -700,9 +767,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.8" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reserved-oxide-symbols" @@ -716,9 +783,9 @@ checksum = "b74b56ffa8bb2830709a538c2cbcae9aa062db0d2a42563bfb09bdaae44020eb" [[package]] name = "rustc-hash" -version = "2.1.1" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustc_apfloat" @@ -744,6 +811,7 @@ dependencies = [ "mir-lower", "once_cell", "oxide-artifacts", + "oxide-tile-export", "pliron", "pliron-derive", "reserved-oxide-symbols", @@ -772,9 +840,9 @@ checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "safe_arch" -version = "0.7.4" +version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" +checksum = "42c6efa15875e6ecb39ca61fb0b0c1a40b84fac5a5ffe71eef7d1000c8eb3f5f" dependencies = [ "bytemuck", ] @@ -812,7 +880,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn", ] [[package]] @@ -850,35 +918,24 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.15.1" +version = "1.16.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +checksum = "5b3dc8af474f516a851ff4bd12db780f948b9250ad37211e4eec0bccea54e01b" [[package]] name = "spin" -version = "0.12.0" +version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1527984ca054dfca79333baec451042863f485fbee01b7bf6d911de915cac865" +checksum = "0134f9043ed38b087ac4f7d4af44c79e2c9e5094421fe3164f435ce585953b10" dependencies = [ "lock_api", ] [[package]] name = "syn" -version = "2.0.114" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "syn" -version = "3.0.3" +version = "3.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" dependencies = [ "proc-macro2", "quote", @@ -900,29 +957,29 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] name = "triple_arena" -version = "0.14.0" +version = "0.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a99051ec05383f70bcab71884d39083a455257795e71cb780048a4bc890f3b8" +checksum = "eee55b64228381514ddc4ad004e631ccf45e594e428d13ce197697ce56ef1c55" dependencies = [ "recasting", ] @@ -962,9 +1019,9 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "wide" -version = "0.7.33" +version = "1.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce5da8ecb62bcd8ec8b7ea19f69a51275e91299be594ea5cc6ef7819e16cd03" +checksum = "d920ac99c3c8edce110cb8d07dbb324d6d026011dce85b1e9355b70f0adacc4f" dependencies = [ "bytemuck", "safe_arch", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/Cargo.toml index dc172e91a1..67b39a8c36 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/Cargo.toml @@ -4,7 +4,7 @@ version = "0.2.1" edition = "2024" authors = ["Nihal Pasham "] license = "Apache-2.0" -repository = "https://github.com/NVlabs/cuda-oxide.git" +repository = "https://github.com/NVIDIA/cuda-rust" description = "Unified host/device compilation backend for rustc (CUDA/PTX codegen)" readme = "README.md" @@ -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] @@ -24,6 +24,8 @@ llvm-export = { path = "../llvm-export" } dialect-mir = { path = "../dialect-mir" } dialect-nvvm = { path = "../dialect-nvvm" } reserved-oxide-symbols = { path = "../reserved-oxide-symbols" } +# `#[tile_export]` descriptions, read back out of the metadata functions' MIR. +oxide-tile-export = { path = "../oxide-tile-export" } oxide-artifacts = { version = "0.2.1", features = ["object-write"] } # Shared driver-independent libNVVM + nvJitLink finalization policy. This # crate loads CUDA compiler libraries at run time and never links libcuda. @@ -31,12 +33,12 @@ cuda-artifact-finalizer = { path = "../cuda-artifact-finalizer" } # IMPORTANT: This crate is isolated from the main workspace, so we cannot use # `{ workspace = true }`. Keep these versions synchronized with root Cargo.toml -# [workspace.dependencies] section. Source MUST match the root (same git + rev), -# or pliron resolves to two distinct crates and types stop unifying. -pliron = { git = "https://github.com/pliron-org/pliron", rev = "06d99a3df938394583715aacf5eadb65b4791c09" } -pliron-derive = { git = "https://github.com/pliron-org/pliron", rev = "06d99a3df938394583715aacf5eadb65b4791c09" } +# [workspace.dependencies] section. Use the same registry source and exact +# release, or pliron resolves to distinct crates and types stop unifying. +pliron = { version = "=0.20.0" } +pliron-derive = { version = "=0.20.0" } linkme = "0.3" # must match workspace -combine = "4.6" # must match workspace +combine = "4.6.8" # must match workspace; Pliron needs SourcePosition: Hash once_cell = "1.18" # must match workspace thiserror = "2" # must match workspace serde_json = "1" # must match workspace diff --git a/cuda-oxide/crates/rustc-codegen-cuda/README.md b/cuda-oxide/crates/rustc-codegen-cuda/README.md index c1a7f2d550..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-08-28`. +The backend uses `#![feature(rustc_private)]` (unlocked on stable via `RUSTC_BOOTSTRAP=1`) and pins to a specific toolchain release via `rust-toolchain.toml`. The workspace currently uses `1.99.1`. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/STATUS.md b/cuda-oxide/crates/rustc-codegen-cuda/STATUS.md index 6cb3a33193..94b339cbda 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/STATUS.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/STATUS.md @@ -25,8 +25,11 @@ Run `cuda-oxide/scripts/check-error-example-status.sh` to verify both are in syn | `error_host_arch_intrinsic` | diagnostics-fixture | `core::arch::` intrinsic reached | | `error_host_target_feature` | diagnostics-fixture | `#[target_feature]` fn reached | | `error_kernel_shared_param` | diagnostics-fixture | AS3/AS5 pointer as kernel parameter | +| `error_kernel_shared_ptr_param` | diagnostics-fixture | AS3 pointer inside a kernel parameter struct | | `error_missing_device_attr` | diagnostics-fixture | `thread::index_*` stub (#76) | | `error_set_discriminant_uninhabited` | diagnostics-fixture | Invalid enum variant selection | +| `error_union_shared_pointer_addr` | diagnostics-fixture | `&mut`/in-place write to a shared-pointer union field | +| `error_union_shared_pointer_layout` | support-gap | Oversized AS3 pointer array in a union field | Drops whose monomorphized glue is provably a no-op (e.g. the `core::array::IntoIter` behind `for x in arr` with Copy elements) lower @@ -37,9 +40,10 @@ other device function; see the `drop_glue` example. This covers direct `impl Drop` types reachable from kernels. The shim and every `Drop::drop` body it reaches go through the same device pipeline, so drop glue whose MIR uses constructs the pipeline cannot yet translate -(e.g. slice/`Vec` element drops, panic formatting) still fails +(e.g. `Vec` element drops, panic formatting) still fails compilation with a diagnostic. Destructors are never silently skipped: -a drop is either proven unobservable or its call is emitted. +a drop is either proven unobservable and terminating, or its call is +emitted. The remaining enum storage fixtures deliberately fail closed. Pointer and integer payloads cannot share one lowered slot without erasing LLVM pointer @@ -54,3 +58,13 @@ Larger expansions remain rejected to keep generated code shape bounded, while pointer vectors still need separate ABI and address-space-cast semantics. (Bool leaves use canonical i8 physical bytes and are converted recursively at the same boundary.) + +The two union fixtures fail closed for related reasons. A union field that +holds a shared-memory pointer, directly or through structs and bounded arrays, +uses the same target-stable generic-pointer storage, converted at the union's +construct and extract boundaries, and the same 16-leaf bound applies to each +field. Taking the in-place address of such a field (`u.ptr = p`, `&u.ptr`, +`&mut u.ptr`, or a write into a struct nested in it) is refused: the loads and +stores made through that address carry the field's semantic type, the raw +shared pointer, which the generic-pointer bytes cannot honor. Build the union +by value and read its fields by value instead. 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 1d51a50e54..894c5bdc74 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/abi_hmm/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/Cargo.toml index 4e2527e09f..4adb6cefb6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/abi_hmm/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 b8e441c59e..aec55dddd1 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/addressof_sharedarray/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/Cargo.toml index 062b38831e..23c939ec01 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/addressof_sharedarray/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 0b75539022..66d6f345c4 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/aligned_field_loads/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/Cargo.toml index bc68d54d7c..36307dc4f7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/Cargo.toml @@ -22,4 +22,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/verify-code-shape.sh index a4724bdfa4..07569156a9 100755 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/verify-code-shape.sh +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_loads/verify-code-shape.sh @@ -5,7 +5,7 @@ # The runtime SUCCESS gate only proves the kernels are bit-correct, which # holds with or without the optimization this example exists to demonstrate. # The perf property lives in the PTX shape: an #[repr(C, align(8))] pair must -# fuse its two adjacent f32 field reads into one ld.global.v2.b32, while the +# fuse its two adjacent f32 field reads into one vector load, while the # natural-align-4 control must keep two scalar loads. Assert both directions # so a silent regression of either the fusion or the control cannot pass. @@ -57,8 +57,10 @@ reject_entry_shape() { fi } -vector_load='ld\.global\.v2\.b32' -scalar_load='ld\.global\.b32' +# Read-only, non-aliasing inputs may use the non-coherent read-only cache +# (`ld.global.nc`). That changes the cache path, not the load width or alignment. +vector_load='ld\.global(\.nc)?\.v2\.b32' +scalar_load='ld\.global(\.nc)?\.b32' # The align(8) kernels must fuse the pair into a single wide load. The `lanes` # kernels reach the same two f32 through an array index rather than named 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 502e06c899..d3a69ad842 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/aligned_field_stores/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/Cargo.toml index b520ce072a..3a527c0ed3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/aligned_field_stores/Cargo.toml @@ -22,4 +22,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 708068528d..2ace543695 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/array_constants/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/Cargo.toml index f5515f5a8c..d1e4c14929 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_constants/Cargo.toml @@ -23,4 +23,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 0c024bb8a5..50e069c917 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/array_for_loop/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/Cargo.toml index 930d5b99a7..f8d9cb3146 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_for_loop/Cargo.toml @@ -24,4 +24,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 e83e8a8704..9aacc03b6a 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/array_index/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/Cargo.toml index d29539a230..074754ac7a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_index/Cargo.toml @@ -21,4 +21,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 b93c70c05f..23873c746e 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/array_init/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/Cargo.toml index 14da9fa8f4..bcd04614b1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/array_init/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 0401eff7d5..11978cf8c3 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/asin_acos_smoke/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/Cargo.toml index 5970b18e0e..c217a7e73a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/asin_acos_smoke/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 248ff82b0f..57c5ac3626 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/Cargo.toml index 8a7fc7adbb..6bc61aad54 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/assume_init_uninhabited_trap/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 22a8b4bc2c..ab1c6b4fec 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 @@ -145,7 +145,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -172,7 +172,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -183,7 +183,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -204,6 +204,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -217,6 +218,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -443,6 +445,12 @@ 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" @@ -558,6 +566,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.18", +] + [[package]] name = "parking_lot_core" version = "0.9.12" @@ -712,6 +729,19 @@ dependencies = [ "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" @@ -794,9 +824,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", @@ -950,3 +980,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/async_mlp/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/Cargo.toml index b364484300..f76bed559e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_mlp/Cargo.toml @@ -15,6 +15,6 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host", features = ["async"] } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } -cuda-async = { path = "../../../../../cuda-async", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } +cuda-async = { path = "../../../../../cuda-async" } tokio = { version = "1", features = ["rt", "rt-multi-thread", "macros"] } 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 902af286c7..3b94cbfb24 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 @@ -145,7 +145,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -172,7 +172,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -183,7 +183,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -204,6 +204,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -217,6 +218,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -443,6 +445,12 @@ 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" @@ -558,6 +566,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.18", +] + [[package]] name = "parking_lot_core" version = "0.9.12" @@ -712,6 +729,19 @@ dependencies = [ "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" @@ -794,9 +824,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", @@ -950,3 +980,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/async_vecadd/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/Cargo.toml index 6fed35118f..aa5cbb40e5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/async_vecadd/Cargo.toml @@ -15,6 +15,6 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host", features = ["async"] } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } -cuda-async = { path = "../../../../../cuda-async", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } +cuda-async = { path = "../../../../../cuda-async" } tokio = { version = "1", features = ["rt", "rt-multi-thread", "macros"] } 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 25e548ea3b..2eedfa7f06 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/atomic_f16/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/Cargo.toml index 23a4ad2699..2fdb444522 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomic_f16/Cargo.toml @@ -25,4 +25,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Primitive `f16` is nightly-only, so the shared crate gates its `DeviceCopy` # impl behind this feature; `half::f16` needs nothing extra. -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0", features = ["f16"] } +cuda-core = { path = "../../../../../cuda-core", features = ["f16"] } 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 ba4ca416b9..ba9f8d35c8 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 @@ -27,37 +27,41 @@ const BINS: &[u32] = &[1, 4, 16, 64, 256, 1024, 4096]; mod kernels { use super::*; + // Atomic storage arrives as a raw writable pointer: a shared &[f16] or + // &[f32] would promise immutable elements even after a cast to an atomic. + // The unsafe launch keeps storage alive and excludes non-atomic access. + #[kernel] - pub fn f32_add(hist: &[f32], n: u32, nbins: u32) { + pub fn f32_add(hist: *mut f32, n: u32, nbins: u32) { let i = thread::index_1d().get(); if i >= n as usize { return; } let bin = i % nbins as usize; - let cell = unsafe { &*(hist.as_ptr().add(bin) as *const DeviceAtomicF32) }; + let cell = unsafe { DeviceAtomicF32::from_ptr(hist.add(bin)) }; cell.fetch_add(1.0f32, AtomicOrdering::Relaxed); } #[kernel] - pub fn f16_add(hist: &[f16], n: u32, nbins: u32) { + pub fn f16_add(hist: *mut f16, n: u32, nbins: u32) { let i = thread::index_1d().get(); if i >= n as usize { return; } let bin = i % nbins as usize; - let cell = unsafe { &*(hist.as_ptr().add(bin) as *const DeviceAtomicF16) }; + let cell = unsafe { DeviceAtomicF16::from_ptr(hist.add(bin)) }; cell.fetch_add(1.0f16, AtomicOrdering::Relaxed); } #[kernel] - pub fn f32_add_return(hist: &[f32], mut old: DisjointSlice, n: u32, nbins: u32) { + pub fn f32_add_return(hist: *mut f32, mut old: DisjointSlice, n: u32, nbins: u32) { let gid = thread::index_1d(); let i = gid.get(); if i >= n as usize { return; } let bin = i % nbins as usize; - let cell = unsafe { &*(hist.as_ptr().add(bin) as *const DeviceAtomicF32) }; + let cell = unsafe { DeviceAtomicF32::from_ptr(hist.add(bin)) }; let prev = cell.fetch_add(1.0f32, AtomicOrdering::Relaxed); if let Some(out) = old.get_mut(gid) { *out = prev; @@ -65,14 +69,14 @@ mod kernels { } #[kernel] - pub fn f16_add_return(hist: &[f16], mut old: DisjointSlice, n: u32, nbins: u32) { + pub fn f16_add_return(hist: *mut f16, mut old: DisjointSlice, n: u32, nbins: u32) { let gid = thread::index_1d(); let i = gid.get(); if i >= n as usize { return; } let bin = i % nbins as usize; - let cell = unsafe { &*(hist.as_ptr().add(bin) as *const DeviceAtomicF16) }; + let cell = unsafe { DeviceAtomicF16::from_ptr(hist.add(bin)) }; let prev = cell.fetch_add(1.0f16, AtomicOrdering::Relaxed); if let Some(out) = old.get_mut(gid) { *out = prev; @@ -80,6 +84,12 @@ mod kernels { } } +// Request writable-buffer authority without creating a Rust slice over the +// device allocation. Callers keep it alive and synchronize before host reads. +fn atomic_storage(buffer: &mut DeviceBuffer) -> *mut T { + buffer.cu_deviceptr() as *mut T +} + fn main() -> Result<(), Box> { let ctx = CudaContext::new(0)?; let stream = ctx.default_stream(); @@ -101,20 +111,22 @@ fn bench_no_return( stream: &Arc, bins: u32, ) -> Result<(), Box> { - let hist_f32 = DeviceBuffer::::zeroed(stream, bins as usize)?; - let hist_f16 = DeviceBuffer::::zeroed(stream, bins as usize)?; + let mut hist_f32 = DeviceBuffer::::zeroed(stream, bins as usize)?; + let mut hist_f16 = DeviceBuffer::::zeroed(stream, bins as usize)?; let cfg = LaunchConfig::for_num_elems(N); let f32_ms = time_gpu_iters(stream, ITERS, || { reset(&hist_f32, stream)?; - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. - unsafe { module.f32_add(stream, cfg, &hist_f32, N, bins) }?; + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. + unsafe { module.f32_add(stream, cfg, atomic_storage(&mut hist_f32), N, bins) }?; Ok(()) })?; let f16_ms = time_gpu_iters(stream, ITERS, || { reset(&hist_f16, stream)?; - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. - unsafe { module.f16_add(stream, cfg, &hist_f16, N, bins) }?; + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. + unsafe { module.f16_add(stream, cfg, atomic_storage(&mut hist_f16), N, bins) }?; Ok(()) })?; @@ -128,22 +140,42 @@ fn bench_return( stream: &Arc, bins: u32, ) -> Result<(), Box> { - let hist_f32 = DeviceBuffer::::zeroed(stream, bins as usize)?; - let hist_f16 = DeviceBuffer::::zeroed(stream, bins as usize)?; + let mut hist_f32 = DeviceBuffer::::zeroed(stream, bins as usize)?; + let mut hist_f16 = DeviceBuffer::::zeroed(stream, bins as usize)?; let mut old_f32 = DeviceBuffer::::zeroed(stream, N as usize)?; let mut old_f16 = DeviceBuffer::::zeroed(stream, N as usize)?; let cfg = LaunchConfig::for_num_elems(N); let f32_ms = time_gpu_iters(stream, ITERS, || { reset(&hist_f32, stream)?; - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. - unsafe { module.f32_add_return(stream, cfg, &hist_f32, &mut old_f32, N, bins) }?; + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. + unsafe { + module.f32_add_return( + stream, + cfg, + atomic_storage(&mut hist_f32), + &mut old_f32, + N, + bins, + ) + }?; Ok(()) })?; let f16_ms = time_gpu_iters(stream, ITERS, || { reset(&hist_f16, stream)?; - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. - unsafe { module.f16_add_return(stream, cfg, &hist_f16, &mut old_f16, N, bins) }?; + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. + unsafe { + module.f16_add_return( + stream, + cfg, + atomic_storage(&mut hist_f16), + &mut old_f16, + N, + bins, + ) + }?; Ok(()) })?; 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 5e00e7c262..94a77623bb 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 @@ -19,8 +19,12 @@ use cuda_host::cuda_module; mod kernels { use super::*; + // Atomic storage arrives as a raw writable pointer: &[f16] would promise + // immutable elements even after a cast to an atomic. + // The unsafe launch keeps storage alive and excludes non-atomic access. + #[kernel] - pub fn device_hist_f16(hist: &[f16], mut old: DisjointSlice, n: u32, nbins: u32) { + pub fn device_hist_f16(hist: *mut f16, mut old: DisjointSlice, n: u32, nbins: u32) { let gid = thread::index_1d(); let i = gid.get(); if i >= n as usize { @@ -28,7 +32,7 @@ mod kernels { } let bin = i % nbins as usize; - let cell = unsafe { &*(hist.as_ptr().add(bin) as *const DeviceAtomicF16) }; + let cell = unsafe { DeviceAtomicF16::from_ptr(hist.add(bin)) }; let prev = cell.fetch_add(1.0f16, AtomicOrdering::Relaxed); if let Some(out_elem) = old.get_mut(gid) { *out_elem = prev; @@ -36,10 +40,10 @@ mod kernels { } #[kernel] - pub fn block_hist_f16(hist: &[f16], mut old: DisjointSlice) { + pub fn block_hist_f16(hist: *mut f16, mut old: DisjointSlice) { let gid = thread::index_1d(); let block = thread::blockIdx_x() as usize; - let cell = unsafe { &*(hist.as_ptr().add(block) as *const BlockAtomicF16) }; + let cell = unsafe { BlockAtomicF16::from_ptr(hist.add(block)) }; let prev = cell.fetch_add(1.0f16, AtomicOrdering::Relaxed); if let Some(out_elem) = old.get_mut(gid) { *out_elem = prev; @@ -47,9 +51,9 @@ mod kernels { } #[kernel] - pub fn device_sub_f16(counter: &[f16], mut old: DisjointSlice) { + pub fn device_sub_f16(counter: *mut f16, mut old: DisjointSlice) { let gid = thread::index_1d(); - let cell = unsafe { &*(counter.as_ptr() as *const DeviceAtomicF16) }; + let cell = unsafe { DeviceAtomicF16::from_ptr(counter) }; let prev = cell.fetch_sub(1.0f16, AtomicOrdering::Relaxed); if let Some(out_elem) = old.get_mut(gid) { *out_elem = prev; @@ -57,9 +61,11 @@ mod kernels { } #[kernel] - pub fn system_add_f16(counter: &[f16], mut old: DisjointSlice) { + pub fn system_add_f16(counter: *mut f16, mut old: DisjointSlice) { let gid = thread::index_1d(); - let cell = unsafe { &*(counter.as_ptr() as *const SystemAtomicF16) }; + // SAFETY: the live, aligned device allocation is accessed atomically + // only by this GPU; system scope does not require concurrent CPU access. + let cell = unsafe { SystemAtomicF16::from_ptr(counter) }; let prev = cell.fetch_add(1.0f16, AtomicOrdering::Relaxed); if let Some(out_elem) = old.get_mut(gid) { *out_elem = prev; @@ -67,13 +73,13 @@ mod kernels { } #[kernel] - pub fn device_swap_f16(cell: &[f16], vals: &[f16], mut old: DisjointSlice, n: u32) { + pub fn device_swap_f16(cell: *mut f16, vals: &[f16], mut old: DisjointSlice, n: u32) { let gid = thread::index_1d(); let i = gid.get(); if i >= n as usize { return; } - let cell = unsafe { &*(cell.as_ptr() as *const DeviceAtomicF16) }; + let cell = unsafe { DeviceAtomicF16::from_ptr(cell) }; let prev = cell.swap(vals[i], AtomicOrdering::Relaxed); if let Some(out_elem) = old.get_mut(gid) { *out_elem = prev; @@ -81,13 +87,13 @@ mod kernels { } #[kernel] - pub fn device_store_load_f16(cells: &[f16], mut out: DisjointSlice, n: u32) { + pub fn device_store_load_f16(cells: *mut f16, mut out: DisjointSlice, n: u32) { let gid = thread::index_1d(); let i = gid.get(); if i >= n as usize { return; } - let cell = unsafe { &*(cells.as_ptr().add(i) as *const DeviceAtomicF16) }; + let cell = unsafe { DeviceAtomicF16::from_ptr(cells.add(i)) }; let seen = cell.load(AtomicOrdering::Acquire); cell.store(seen + 1.0f16, AtomicOrdering::Release); if let Some(out_elem) = out.get_mut(gid) { @@ -96,6 +102,12 @@ mod kernels { } } +// Request writable-buffer authority without creating a Rust slice over the +// device allocation. Callers keep it alive and synchronize before host reads. +fn atomic_storage(buffer: &mut DeviceBuffer) -> *mut T { + buffer.cu_deviceptr() as *mut T +} + fn main() { println!("=== DeviceAtomicF16 / BlockAtomicF16 tests ===\n"); @@ -122,14 +134,15 @@ fn run_device_hist( ) { println!("--- DeviceAtomicF16 histogram: n={n}, bins={nbins} ---"); - let hist = DeviceBuffer::::zeroed(stream, nbins as usize).unwrap(); + let mut hist = DeviceBuffer::::zeroed(stream, nbins as usize).unwrap(); let mut old = DeviceBuffer::::zeroed(stream, n as usize).unwrap(); - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. unsafe { module.device_hist_f16( stream, LaunchConfig::for_num_elems(n), - &hist, + atomic_storage(&mut hist), &mut old, n, nbins, @@ -158,12 +171,20 @@ fn run_device_sub(module: &kernels::LoadedModule, stream: &cuda_core::CudaStream println!("--- DeviceAtomicF16 fetch_sub: n={n} ---"); let initial = vec![n as f16]; - let counter = DeviceBuffer::from_host(stream, &initial).unwrap(); + let mut counter = DeviceBuffer::from_host(stream, &initial).unwrap(); let mut old = DeviceBuffer::::zeroed(stream, n as usize).unwrap(); - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. - unsafe { module.device_sub_f16(stream, LaunchConfig::for_num_elems(n), &counter, &mut old) } - .expect("Kernel launch failed"); + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. + unsafe { + module.device_sub_f16( + stream, + LaunchConfig::for_num_elems(n), + atomic_storage(&mut counter), + &mut old, + ) + } + .expect("Kernel launch failed"); stream.synchronize().unwrap(); @@ -180,12 +201,20 @@ fn run_device_sub(module: &kernels::LoadedModule, stream: &cuda_core::CudaStream fn run_system_add(module: &kernels::LoadedModule, stream: &cuda_core::CudaStream, n: u32) { println!("--- SystemAtomicF16 fetch_add: n={n} ---"); - let counter = DeviceBuffer::::zeroed(stream, 1).unwrap(); + let mut counter = DeviceBuffer::::zeroed(stream, 1).unwrap(); let mut old = DeviceBuffer::::zeroed(stream, n as usize).unwrap(); - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. - unsafe { module.system_add_f16(stream, LaunchConfig::for_num_elems(n), &counter, &mut old) } - .expect("Kernel launch failed"); + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. + unsafe { + module.system_add_f16( + stream, + LaunchConfig::for_num_elems(n), + atomic_storage(&mut counter), + &mut old, + ) + } + .expect("Kernel launch failed"); stream.synchronize().unwrap(); @@ -205,17 +234,18 @@ fn run_system_add(module: &kernels::LoadedModule, stream: &cuda_core::CudaStream fn run_device_swap(module: &kernels::LoadedModule, stream: &cuda_core::CudaStream, n: u32) { println!("--- DeviceAtomicF16 swap: n={n} ---"); - let cell = DeviceBuffer::::zeroed(stream, 1).unwrap(); + let mut cell = DeviceBuffer::::zeroed(stream, 1).unwrap(); let host_vals: Vec = (1..=n).map(|v| v as f16).collect(); let vals = DeviceBuffer::from_host(stream, &host_vals).unwrap(); let mut old = DeviceBuffer::::zeroed(stream, n as usize).unwrap(); - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. unsafe { module.device_swap_f16( stream, LaunchConfig::for_num_elems(n), - &cell, + atomic_storage(&mut cell), &vals, &mut old, n, @@ -249,12 +279,19 @@ fn run_device_store_load(module: &kernels::LoadedModule, stream: &cuda_core::Cud println!("--- DeviceAtomicF16 load/store: n={n} ---"); let host_init: Vec = (0..n).map(|v| v as f16).collect(); - let cells = DeviceBuffer::from_host(stream, &host_init).unwrap(); + let mut cells = DeviceBuffer::from_host(stream, &host_init).unwrap(); let mut out = DeviceBuffer::::zeroed(stream, n as usize).unwrap(); - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. unsafe { - module.device_store_load_f16(stream, LaunchConfig::for_num_elems(n), &cells, &mut out, n) + module.device_store_load_f16( + stream, + LaunchConfig::for_num_elems(n), + atomic_storage(&mut cells), + &mut out, + n, + ) } .expect("Kernel launch failed"); @@ -289,7 +326,7 @@ fn run_block_hist( "--- BlockAtomicF16 histogram: blocks={blocks}, threads/block={threads_per_block} ---" ); - let hist = DeviceBuffer::::zeroed(stream, blocks as usize).unwrap(); + let mut hist = DeviceBuffer::::zeroed(stream, blocks as usize).unwrap(); let mut old = DeviceBuffer::::zeroed(stream, (blocks * threads_per_block) as usize).unwrap(); let cfg = LaunchConfig { @@ -298,8 +335,10 @@ fn run_block_hist( shared_mem_bytes: 0, }; - // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. - unsafe { module.block_hist_f16(stream, cfg, &hist, &mut old) }.expect("Kernel launch failed"); + // SAFETY: buffers cover all accesses and stay alive until synchronization. + // Atomic storage has no concurrent non-atomic accesses. + unsafe { module.block_hist_f16(stream, cfg, atomic_storage(&mut hist), &mut old) } + .expect("Kernel launch failed"); stream.synchronize().unwrap(); 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 011e726425..e62f2b5116 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.lock @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/atomics/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.toml index 7b389ad57e..c6a8d91cb5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/atomics/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 2d839810ca..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-08-28`) 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-08-28` 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..ddaca5cd7a 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 @@ -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 f17829532f..167540b149 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.lock @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/barrier/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.toml index 31e593b19d..66bb64ee0f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/barrier/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 31c36820b6..21784a78d9 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/bf16x2_arith/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/Cargo.toml index 71c2e075b9..d48988e830 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_arith/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.lock index 4edfd7f252..02c07a415e 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.toml index 844f43828e..b673d9f997 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bf16x2_fma/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.lock index 876fb8c273..dd5b1d27f2 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.toml index 04998f1938..dd629b7980 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/bool_phi_cmp/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 c41a3f5c30..4029312b7b 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/carrying_mul_add/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/Cargo.toml index 26e380c743..b91be22980 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/carrying_mul_add/Cargo.toml @@ -20,4 +20,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 0476d620a1..23b2870e9f 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cast_chunks/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/Cargo.toml index 4e9d2be127..5fd3cef2d0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_chunks/Cargo.toml @@ -21,4 +21,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 a885f34945..3b2f05431c 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cast_tests/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/Cargo.toml index 75faf1a877..da6e45546e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cast_tests/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 6250c3b0b7..ed4913c237 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cbrt_smoke/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/Cargo.toml index f12349e649..0150e85e83 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cbrt_smoke/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 22a605d5c9..68c6a15821 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/checked_arith/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/Cargo.toml index 2ca5d95cbe..ff745df1fe 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/checked_arith/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 6409d9f788..e47f313cfc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.lock @@ -138,7 +138,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,7 +150,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -194,6 +195,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/clc/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.toml index e86891a1fc..fe9ed90e41 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/clc/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-macros = { path = "../../../../crates/cuda-macros" } cuda-device = { path = "../../../../crates/cuda-device" } cuda-host = { path = "../../../../crates/cuda-host" } 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 53a895695a..5d785dbfe8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.lock @@ -138,7 +138,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,7 +150,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -194,6 +195,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cluster/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.toml index 0d94935130..f3b6d17891 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-macros = { path = "../../../../crates/cuda-macros" } cuda-device = { path = "../../../../crates/cuda-device" } cuda-host = { path = "../../../../crates/cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/README.md index e87b4f8ca4..6e543ac139 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/README.md @@ -9,7 +9,7 @@ Demonstrates Thread Block Clusters, a Hopper feature that enables direct shared 1. **test_cluster_compile_time**: Compile-time cluster configuration with `#[cluster_launch]` 2. **test_cluster_intrinsics**: Cluster special registers (ctaid, nctaid, rank, size) 3. **test_cluster_sync**: Cluster-wide synchronization -4. **test_dsmem_ring_exchange**: Distributed shared memory read through `map_shared_rank` + dereference +4. **test_dsmem_ring_exchange**: Distributed shared memory read through `map_shared_rank` + dereference, with the mapped pointer also round-tripped through a `(ptr, rank)` tuple and a device fn (`bounce_cluster_pair`) and compared against the direct read 5. **test_dsmem_reduction**: All-to-one reduction using `dsmem_read_u32` 6. **test_dsmem_mapped_store**: Distributed shared memory write through `map_shared_rank_mut` + dereference 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 0e78973abd..d9d2a5c18f 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 @@ -15,7 +15,7 @@ //! Uses unified compilation: single `cargo oxide run cluster` use core::ptr::{addr_of, addr_of_mut}; -use cuda_device::{DisjointSlice, SharedArray, cluster, cluster_launch, kernel, thread}; +use cuda_device::{DisjointSlice, SharedArray, cluster, cluster_launch, device, kernel, thread}; use cuda_host::cuda_module; // ============================================================================ @@ -136,6 +136,12 @@ mod kernels { // Test 3: Distributed Shared Memory (Ring Exchange) // ============================================================================ + #[inline(never)] + #[device] + fn bounce_cluster_pair(pair: (*const u32, u64)) -> (*const u32, u64) { + pair + } + /// Test kernel for distributed shared memory via map_shared_rank() + dereference. #[kernel] #[cluster_launch(4, 1, 1)] @@ -160,7 +166,16 @@ mod kernels { let neighbor_rank = (my_rank + 1) % cluster_size; let neighbor_ptr = unsafe { cluster::map_shared_rank(addr_of!(SHMEM) as *const u32, neighbor_rank) }; - let neighbor_value = unsafe { *neighbor_ptr }; + let pair = bounce_cluster_pair((neighbor_ptr, neighbor_rank as u64)); + + let direct = unsafe { *neighbor_ptr }; + let bounced = unsafe { *pair.0 }; + + let neighbor_value = if bounced == direct && pair.1 == neighbor_rank as u64 { + bounced + } else { + u32::MAX + }; let idx = my_rank as usize; if idx < output.len() { diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/verify-code-shape.sh index dd2efc8080..fb56fd21f2 100755 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/verify-code-shape.sh +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cluster/verify-code-shape.sh @@ -71,4 +71,14 @@ require_entry_shape test_dsmem_ring_exchange \ require_entry_shape test_dsmem_mapped_store \ "cluster-shared remote store" 'st\.shared::cluster\.' +# The private #[inline(never)] helper reaches LLVM as `noinline`, but opt -O2 +# still removes the call: the helper returns its only argument unchanged, so +# opt marks that argument `returned` and forwards it at the call site. This +# check therefore reads the pre-opt cluster.ll; the optimized PTX does not +# retain the call. +require_ll_shape "tuple round-trip LLVM call" \ + 'call \{ ptr, i64 \} @bounce_cluster_pair\(' +require_ll_shape "cluster-shared pointer normalized before tuple construction" \ + 'addrspacecast ptr addrspace\(7\) %[A-Za-z0-9_.]+ to ptr' + echo "cluster code shape: PASS" 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 70deb8c7c7..f74d2527f8 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 @@ -138,7 +138,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,7 +150,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -194,6 +195,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/compiler_features/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/Cargo.toml index 502069fdd7..2364fe00c7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/compiler_features/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-macros = { path = "../../../../crates/cuda-macros" } cuda-device = { path = "../../../../crates/cuda-device" } cuda-host = { path = "../../../../crates/cuda-host" } 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 60eb1baa31..8e9267dce6 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/complex_mul/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/Cargo.toml index dd33c6c404..ed7bef938e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/complex_mul/Cargo.toml @@ -14,4 +14,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 3ef71976cc..1b55fe2102 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/const_aggregate/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/Cargo.toml index 962c1b4d29..240de2ad4f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_aggregate/Cargo.toml @@ -19,4 +19,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 41fabb1728..1624a6b9f9 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/Cargo.toml index ccfe0f0cfb..c547025eb8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_bool_dead_branch/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 ccc08f96f2..59af5b4d4c 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 @@ -138,7 +138,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,7 +150,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -194,6 +195,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/const_generic/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/Cargo.toml index 5fe9cb8235..e13cc14616 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_generic/Cargo.toml @@ -10,5 +10,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } ptx-parse = { path = "../../../ptx-parse" } 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 26fb53caca..6378ac770e 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/const_table_lookup/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/Cargo.toml index 10dab4b09b..8c077f80c1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_table_lookup/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 12a18ab687..e5bb7b77b8 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/Cargo.toml index 771f29995e..1d54a395c5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/const_tuple_table_lookup/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.lock index 4470121240..4a5ddfcbfc 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.toml index a6e74dc1d7..f94c5bb3a5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/README.md index 7ffccca3a1..37c174f27a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_index_from_end/README.md @@ -28,8 +28,10 @@ The length must come from the fat value itself, never from the data pointer's pointee type. For `&[[u32; 3]]`, the pointee array length 3 is the row width; the from-end index selects a row from the outer runtime-length slice, so subtracting from 3 would read the wrong row. -From-end indexes without that fat-deref provenance (e.g. after a -`Subslice`) still fail closed instead of guessing a length. +A from-end index after a slice `Subslice` takes the rebuilt subslice +length (the `subslice_projection` example covers that case). Other +bases without fat-deref provenance still fail closed instead of +guessing a length. ## Kernels 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 6e8890fbf5..ec1b311b82 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/constant_memory/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/Cargo.toml index a3a07bcb41..56c035b42a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 ac248d8bb3..254f6e1446 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/Cargo.toml index 4066b3e3b8..614d2601df 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_coeffs/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 98e2a8c034..0290d6b02d 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/constant_memory_simple/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/Cargo.toml index 34e520e814..a68bb22778 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/verify-debug-info.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/verify-debug-info.sh index 69220ef32b..49a38b6182 100755 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/verify-debug-info.sh +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_simple/verify-debug-info.sh @@ -80,7 +80,7 @@ grep -Fq 'expr: !DIExpression())' <<<"${scale_expression}" \ scale_expression_id="$(sed -nE 's/^!([0-9]+) = .*/\1/p' <<<"${scale_expression}")" physical="$(unique_line "@${scale_linkage} = ")" -grep -Eq "^@${scale_linkage} = addrspace\\(4\\) global \\[4 x i8\\] .*align 4, !dbg !${scale_expression_id}$" \ +grep -Eq "^@${scale_linkage} = addrspace\\(4\\) externally_initialized global \\[4 x i8\\] .*align 4, !dbg !${scale_expression_id}$" \ <<<"${physical}" || fail "SCALE physical AS4 storage is not attached to its expression" # Follow the exact semantic type graph. The physical [4 x i8] backing must be 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 7c62ceae52..5ee35f9f49 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/constant_memory_table/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/Cargo.toml index 3621f16401..7ef177c4a1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/constant_memory_table/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 dcca83bb32..1929b90d68 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/coop_groups_demo/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/Cargo.toml index e6f02721e4..c656535920 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/coop_groups_demo/Cargo.toml @@ -21,4 +21,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 fdaba39cbd..861cd541fa 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/Cargo.toml index bc6368973e..0001721a40 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_aggregate_borrow/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 e583f5fddb..e0ad2e575b 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/Cargo.toml index 5d1368dc34..9bc2994651 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/copy_nonoverlapping/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.lock index d6e529758d..670a99a5a0 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.toml index eef635f2ef..aece4f4395 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/counted_barrier/Cargo.toml @@ -11,4 +11,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 0641ecb7ab..e2b60b35e4 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cp_async_small/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/Cargo.toml index 452950d5df..9374daf7d1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_small/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.lock index 19b846ea77..554696b9a9 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 @@ -137,7 +137,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -180,6 +180,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -193,6 +194,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.toml index 7854f32e78..dfce777cf8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cp_async_zfill/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cpp_consumes_rust_device/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cpp_consumes_rust_device/Cargo.lock index 059f389beb..65b1237134 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 @@ -136,7 +136,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -148,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -159,7 +159,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -179,6 +179,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -192,6 +193,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -310,6 +312,12 @@ 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" @@ -416,6 +424,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -530,6 +547,19 @@ dependencies = [ "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" @@ -715,3 +745,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.lock index 5d37f456da..86fd1896bc 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 @@ -138,7 +138,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,7 +150,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -194,6 +195,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -321,6 +323,12 @@ 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" @@ -427,6 +435,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -541,6 +558,19 @@ dependencies = [ "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" @@ -726,3 +756,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.toml index 60e211e141..f13e386aa1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/Cargo.toml @@ -17,4 +17,4 @@ members = ["embedded-kernel-lib"] embedded-kernel-lib = { path = "embedded-kernel-lib" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/embedded-kernel-lib/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/embedded-kernel-lib/Cargo.toml index f7e5437906..89c4720eb4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/embedded-kernel-lib/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_embedded/embedded-kernel-lib/Cargo.toml @@ -6,6 +6,6 @@ authors = ["Nihal Pasham "] license = "Apache-2.0" [dependencies] -cuda-core = { path = "../../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../../cuda-core" } cuda-device = { path = "../../../../cuda-device" } cuda-host = { path = "../../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.lock index 8aa5c60c45..93b3a02f35 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 @@ -139,7 +139,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -162,7 +162,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -182,6 +182,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -195,6 +196,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -313,6 +315,12 @@ dependencies = [ "either", ] +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + [[package]] name = "kernel-lib" version = "0.1.0" @@ -428,6 +436,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -542,6 +559,19 @@ dependencies = [ "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" @@ -727,3 +757,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.toml index 6f6aef3b79..4a75520499 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/Cargo.toml @@ -15,5 +15,5 @@ kernel-lib = { path = "kernel-lib" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } ptx-parse = { path = "../../../ptx-parse" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/kernel-lib/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/kernel-lib/Cargo.toml index ab0b7bf067..3300772048 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/kernel-lib/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_kernel/kernel-lib/Cargo.toml @@ -6,6 +6,6 @@ authors = ["Nihal Pasham "] license = "Apache-2.0" [dependencies] -cuda-core = { path = "../../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../../cuda-core" } cuda-device = { path = "../../../../cuda-device" } cuda-host = { path = "../../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.lock index 7d6afea4d1..65317d8b2c 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 @@ -138,7 +138,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,7 +150,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -194,6 +195,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -427,6 +435,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -541,6 +558,19 @@ dependencies = [ "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" @@ -726,3 +756,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.toml index 1061dcaa2e..c6454380d0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/Cargo.toml @@ -15,4 +15,4 @@ members = ["merged-symbols-lib"] merged-symbols-lib = { path = "merged-symbols-lib" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/merged-symbols-lib/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/merged-symbols-lib/Cargo.toml index 87ded48a34..0917964f07 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/merged-symbols-lib/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_merged_symbols/merged-symbols-lib/Cargo.toml @@ -5,6 +5,6 @@ edition = "2024" license = "Apache-2.0" [dependencies] -cuda-core = { path = "../../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../../cuda-core" } cuda-device = { path = "../../../../cuda-device" } cuda-host = { path = "../../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.lock index 304d4f8738..02fa70fa98 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 @@ -138,7 +138,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -150,7 +150,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -181,6 +181,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -194,6 +195,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -536,6 +553,19 @@ dependencies = [ "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" @@ -721,3 +751,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.toml index 3884576b07..28b3b98944 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cross_crate_recursion/Cargo.toml @@ -13,4 +13,4 @@ members = ["reprolib"] reprolib = { path = "reprolib" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 b7ccbf3cd6..4265107f33 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cuda_module_contract/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_contract/Cargo.toml index 7b84833f1d..f0254eed86 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_contract/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_contract/Cargo.toml @@ -10,5 +10,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } ptx-parse = { path = "../../../ptx-parse" } 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 aedf8ee1f5..e21415db54 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -427,6 +435,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -541,6 +558,19 @@ dependencies = [ "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" @@ -726,3 +756,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/Cargo.toml index 60a96839c2..50057bb834 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/Cargo.toml @@ -15,4 +15,4 @@ module-kernels = { path = "kernel-lib" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/kernel-lib/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/kernel-lib/Cargo.toml index ef1eea5067..4054964789 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/kernel-lib/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_in_lib/kernel-lib/Cargo.toml @@ -6,6 +6,6 @@ authors = ["Nihal Pasham "] license = "Apache-2.0" [dependencies] -cuda-core = { path = "../../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../../cuda-core" } cuda-device = { path = "../../../../cuda-device" } cuda-host = { path = "../../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.lock index 15fa7deaa9..9cffb5991d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.toml index d32bf55c20..25e1f7087b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cuda_module_nested/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 79a4f6425a..a0997177a1 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 @@ -169,7 +169,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -184,7 +184,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -196,7 +196,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -207,7 +207,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -216,7 +216,7 @@ dependencies = [ [[package]] name = "cutile" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-async", @@ -238,7 +238,7 @@ dependencies = [ [[package]] name = "cutile-compiler" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-async", @@ -258,7 +258,7 @@ dependencies = [ [[package]] name = "cutile-ir" -version = "0.4.0" +version = "0.5.0" dependencies = [ "half", "indexmap", @@ -268,7 +268,7 @@ dependencies = [ [[package]] name = "cutile-macro" -version = "0.4.0" +version = "0.5.0" dependencies = [ "convert_case", "cutile-compiler", @@ -905,9 +905,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.150" +version = "1.0.154" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" dependencies = [ "itoa", "memchr", @@ -1017,9 +1017,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", @@ -1403,6 +1403,6 @@ dependencies = [ [[package]] name = "zmij" -version = "1.0.21" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" 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 3339cb1874..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" } -cuda-async = { path = "../../../../../cuda-async", version = "0.4.0" } -cuda-core = { path = "../../../../../cuda-core", 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/simt/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.lock index b6bb84c873..5b1b40426e 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.toml index 3fc8f8b106..8c3a514759 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cutile_inter_kernel/simt/Cargo.toml @@ -11,4 +11,4 @@ resolver = "2" [dependencies] cuda-device = { path = "../../../../cuda-device" } cuda-host = { path = "../../../../cuda-host" } -cuda-core = { path = "../../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../../cuda-core" } 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 41761b8893..310b4cd386 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cvt_f16x2/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/Cargo.toml index e00aa1988b..b1f2b7609e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_f16x2/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 be34c2de32..a26fc79858 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/cvt_packed/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/Cargo.toml index 9cf328c766..34ac59eaa7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/cvt_packed/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 f66d568a2c..896adc4fe1 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/debug-tests/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/Cargo.toml index c75395b550..534fd6811b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug-tests/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-macros = { path = "../../../../crates/cuda-macros" } cuda-device = { path = "../../../../crates/cuda-device" } cuda-host = { path = "../../../../crates/cuda-host" } 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 c45afc11ea..a9b382b282 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/debug/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.toml index 91d4c45692..acedb17d53 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-macros = { path = "../../../../crates/cuda-macros" } cuda-device = { path = "../../../../crates/cuda-device" } cuda-host = { path = "../../../../crates/cuda-host" } 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 ff0ba980c8..c6fcff1155 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.toml index cb50a56f20..f6f209f655 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/debug_pointer_locals/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.lock index fbfc3fe9d0..cd74327ba3 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/device_closures/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.toml index 94060ed5f9..fe6c047882 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_closures/Cargo.toml @@ -10,5 +10,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-host = { path = "../../../cuda-host" } 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 f9a370b1fd..547e2d265c 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/device_ffi_test/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/Cargo.toml index d0230e2cbd..28ace2a1d4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_ffi_test/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 d9f51af73f..cb70f208f9 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/device_global/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/Cargo.toml index d6ee91b1c1..02e24121a8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_global/Cargo.toml @@ -15,4 +15,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 abbccc4d8d..940b241f11 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/device_panic_trap/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/Cargo.toml index 6d759f3025..5cffceadf7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/device_panic_trap/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 99e58cf0fa..047a387564 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/Cargo.toml index 9db326fab9..66b67cb17f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_from_raw_parts/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.lock index b57449f344..6489cf91da 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.toml index a8af263a6e..f74c51dc53 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/disjoint_slice_len/Cargo.toml @@ -11,7 +11,7 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } # The mir-importer intercepts `DisjointSlice::len` only when rustc's MIR # inliner leaves the call intact; at default release settings the call is 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 927e35e5b9..476fdafb8c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/dotprod/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.toml index 3bd36266b9..b227e7f508 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dotprod/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 3986154648..a813d9446e 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/Cargo.toml index 9055cb7618..2aabc1ec9d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dpx_minmax_chains/Cargo.toml @@ -16,4 +16,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 092d9bad63..f5e4bf7171 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/drop_glue/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/Cargo.toml index 0d962f2f9a..f6bb6b8327 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/Cargo.toml @@ -1,3 +1,5 @@ +cargo-features = ["profile-rustflags"] + [package] name = "drop_glue" version = "0.1.0" @@ -8,6 +10,21 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } + +# The compile-only probes need each destructor's own MIR to reach the importer: +# a CFG cycle in `LoopingDrop::drop`, the recursive call in `RecursiveDrop::drop` +# and the two-way branch in `DiamondDrop::drop`. At default release settings +# rustc's MIR inliner folds the empty callees into the destructors and the +# destructors into their glue first. `#[inline(never)]` is not the answer: +# it is carried to LLVM as `noinline`, so the destructors would stay real PTX +# functions and the optimized kernels would call them instead of showing the +# store-then-self-loop that verify-code-shape.sh checks. Switching MIR inlining +# off for this one crate keeps the calls for LLVM to inline. Scoping the flag +# to this package keeps the shared dependency tree identical to every other +# example; a global RUSTFLAGS entry would re-key every dependency unit and +# force a full second dependency-tree build in CI. +[profile.release.package.drop_glue] +rustflags = ["-Zinline-mir=no"] diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/README.md index bbbb15d6a9..1f2777c8f2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/README.md @@ -42,6 +42,37 @@ distinguish a destructor that ran from one that was correctly suppressed. | 5 | `MaybeUninit::write` + `assume_init_drop` | `0xDEADBEEF` | | 6 | `MaybeUninit::assume_init_read` | `0xDEADBEEF` | +## Compile-only drop proof probes + +Three additional kernels guard the proof that lets the compiler erase a +destructor call: two against erasing a non-terminating destructor (#1360), +one against rejecting a destructor the proof should accept. + +- `diverging_cfg_drop_probe`: the destructor contains `loop {}`. +- `diverging_recursive_drop_probe`: the destructor calls an unconditionally + recursive helper. +- `reconverging_drop_probe`: the destructor is an effect-free `if`/`else` + whose branches join back together. The proof must still accept it, so + the drop lowers to a plain branch. + +Each probe writes a sentinel before the drop and a different value after +it. `verify-code-shape.sh` first asserts the importer's decision in the +pre-opt `drop_glue.ll`: the two diverging probes keep a call to their drop +glue in the kernel body, and the reconverging probe has none. It then +checks the optimized PTX. For the diverging probes the pre-drop store +remains and is immediately followed by an unconditional self-loop, while +the unreachable post-drop store is absent. For the reconverging probe both +stores remain. The PTX adjacency holds because `opt -O2` inlines the +retained glue and the destructors. They carry no inline attribute, since +`#[inline(never)]` reaches LLVM as `noinline` and would keep them real PTX +functions; MIR inlining is off for this crate instead (`Cargo.toml`), so the +importer still sees each destructor's own MIR. + +These kernels are compiled but **never launched by `main`**. Do not launch +the diverging probes manually: their destructors intentionally never +return. The existing runtime checks above continue to exercise terminating +destructors only. + ## Usage Optimized: @@ -56,6 +87,21 @@ Low MIR optimization: CUDA_OXIDE_NO_OPT=1 cargo oxide run drop_glue ``` +Optimized compile-only code checks, from the cuda-oxide repository root: + +```bash +cargo oxide build drop_glue +bash crates/rustc-codegen-cuda/examples/drop_glue/verify-code-shape.sh +``` + +The shape check targets optimized output; rebuild without +`CUDA_OXIDE_NO_OPT=1` before running it. The existing smoketest runner also +discovers and runs this script automatically: + +```bash +scripts/smoketest.sh -o '^drop_glue$' +``` + ## Expected output ```text 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 3e3f455599..316fd51308 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 @@ -41,6 +41,57 @@ impl Drop for DropMarker { } } +// Compile-only regression probes. These destructors intentionally never return; +// their kernels are generated but never launched by `main`. MIR inlining is +// off for this crate (see Cargo.toml) so each destructor reaches the importer +// with its own MIR. +#[allow(dead_code)] +struct LoopingDrop(u8); + +impl Drop for LoopingDrop { + #[allow(clippy::empty_loop)] + fn drop(&mut self) { + loop {} + } +} + +#[allow(unconditional_recursion)] +fn recursive_noop() { + recursive_noop(); +} + +#[allow(dead_code)] +struct RecursiveDrop(u8); + +impl Drop for RecursiveDrop { + fn drop(&mut self) { + recursive_noop(); + } +} + +// Precision control for the proof. An effect-free destructor whose two +// branches join back together returns on every path, so the proof must keep +// accepting it and lower the drop to a plain branch. The callees have visible +// empty bodies. MIR inlining is off for this crate (see Cargo.toml), which +// keeps the diamond in the destructor's MIR; `#[inline(never)]` would reach +// LLVM as `noinline` and keep these helpers real PTX functions. +fn left_noop() {} + +fn right_noop() {} + +#[allow(dead_code)] +struct DiamondDrop(u8); + +impl Drop for DiamondDrop { + fn drop(&mut self) { + if self.0 == 0 { + left_noop() + } else { + right_noop() + } + } +} + #[cuda_module] mod kernels { use super::*; @@ -57,6 +108,71 @@ mod kernels { } } + /// Compile-only regression: a reachable CFG cycle inside Drop::drop + /// must not be classified as an erasable no-op. + #[kernel] + pub fn diverging_cfg_drop_probe(mut out: DisjointSlice) { + if thread::index_1d().get() != 0 { + return; + } + + unsafe { + let slot = out.as_mut_ptr(); + slot.write(0x1111_1111); + + { + let _guard = LoopingDrop(0); + } + + // Semantically unreachable. + slot.write(0x2222_2222); + } + } + + /// Compile-only regression: recursive nontermination in a callee must + /// likewise prevent drop elision. + #[kernel] + pub fn diverging_recursive_drop_probe(mut out: DisjointSlice) { + if thread::index_1d().get() != 0 { + return; + } + + unsafe { + let slot = out.as_mut_ptr(); + slot.write(0x3333_3333); + + { + let _guard = RecursiveDrop(0); + } + + // Semantically unreachable. + slot.write(0x4444_4444); + } + } + + /// Compile-only precision control: an effect-free `if`/`else` inside + /// `Drop::drop` joins back together, so the proof must still erase the + /// drop. The post-drop store goes to a second slot so neither store is + /// dead and both stay visible in the optimized PTX. + #[kernel] + pub fn reconverging_drop_probe(mut out: DisjointSlice) { + if thread::index_1d().get() != 0 { + return; + } + + unsafe { + let slot = out.as_mut_ptr(); + slot.write(0x5555_5555); + + { + let _guard = DiamondDrop(thread::index_1d().get() as u8); + } + + // Reachable: the destructor returns. + slot.add(1).write(0x6666_6666); + } + } + #[kernel] pub fn drop_control_kernel(mut out: DisjointSlice) { if thread::index_1d().get() != 0 { diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/verify-code-shape.sh new file mode 100644 index 0000000000..beb07d83d7 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/drop_glue/verify-code-shape.sh @@ -0,0 +1,156 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# The device-side drop no-op proof may erase a destructor call only when the +# glue is effect-free and returns. Assert the importer's decision in the +# pre-opt drop_glue.ll: the two diverging probes keep a call to their drop +# glue in the kernel body, and the reconverging control has none. Then check +# what LLVM made of it in the optimized PTX, so a regression in either +# direction cannot pass compile-only CI. + +set -euo pipefail + +root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ll="${1:-${root}/drop_glue.ll}" +ptx="${2:-${root}/drop_glue.ptx}" + +for artifact in "${ll}" "${ptx}"; do + if [[ ! -s "${artifact}" ]]; then + echo "error: missing ${artifact}; run cargo oxide build drop_glue without CUDA_OXIDE_NO_OPT" >&2 + exit 1 + fi +done + +python3 - "${ll}" "${ptx}" <<'VERIFY_PY' +import re +import sys + +ll_path, ptx_path = sys.argv[1], sys.argv[2] +ll = open(ll_path, encoding="utf-8").read() +ptx = open(ptx_path, encoding="utf-8").read() +# Strip comments so a mention in a comment cannot satisfy or hide a check. +ll = re.sub(r";[^\n]*", "", ll) +ptx = re.sub(r"/\*.*?\*/|//[^\n]*", "", ptx, flags=re.S) + +# A call to the monomorphized `core::ptr::drop_glue::` shim, in the v0 +# mangling the device build uses. +GLUE_CALL = re.compile(r"call void @\S*4core3ptr9drop_glue\S*\(") + +def kernel_body(name): + header = re.search( + r"^define ptx_kernel void @" + re.escape(name) + r"\(.*$", + ll, + flags=re.M, + ) + if header is None: + raise SystemExit(f"error: missing LLVM kernel {name} in {ll_path}") + + end = re.compile(r"^\}$", flags=re.M).search(ll, header.end()) + if end is None: + raise SystemExit(f"error: unterminated LLVM kernel {name} in {ll_path}") + + return ll[header.end():end.start()] + +def require_glue_call(body, type_name, name): + pattern = rf"call void @\S*drop_glue\d+{type_name}\S*\(" + if re.search(pattern, body) is None: + raise SystemExit( + f"error: {name} lost the drop_glue::<{type_name}> call in {ll_path}" + ) + +def reject_glue_call(body, name): + call = GLUE_CALL.search(body) + if call is not None: + raise SystemExit( + f"error: {name} kept a drop glue call the no-op proof should have " + f"erased: {call.group(0)}" + ) + +def entry_body(name): + header = re.search( + r"\.entry\s+" + re.escape(name) + r"\s*\([^)]*\)\s*\{", + ptx, + ) + if header is None: + raise SystemExit(f"error: missing PTX entry {name}") + + depth = 1 + end = header.end() + + while depth and end < len(ptx): + if ptx[end] == "{": + depth += 1 + elif ptx[end] == "}": + depth -= 1 + end += 1 + + if depth: + raise SystemExit(f"error: unterminated PTX entry {name}") + + return ptx[header.end():end - 1] + +def require_store(body, value, name, label="pre-drop"): + pattern = rf"^\s*st\.global\.b32\s+\[[^\]]+\],\s*{value}\s*;" + store = re.search(pattern, body, flags=re.M) + if store is None: + raise SystemExit( + f"error: {name} lost required {label} store {value}" + ) + return store.end() + +# The PTX adjacency below holds because `opt -O2` inlines the retained glue +# and the destructor into the kernel. The destructors carry no inline +# attribute: `#[inline(never)]` reaches LLVM as `noinline` and would keep them +# real PTX functions, so the kernel would call them instead. MIR inlining is +# off for this crate (Cargo.toml) so the importer still sees each destructor's +# own MIR. The drop_glue.ll assertion is the one that tracks the importer's +# own decision; this check confirms what LLVM made of the call it kept. +def require_nonreturning_loop(body, after_store, name): + # Both optimized probes lower to a store followed by an unconditional + # self-loop. Tie the loop to that store's continuation: an unrelated + # loop elsewhere in the entry would not prove the drop was preserved. + pattern = r"\s*([\w$.]+):\s*bra(?:\.uni)?\s+\1\s*;" + if not re.match(pattern, body[after_store:]): + raise SystemExit( + f"error: {name} pre-drop store is not followed by an unconditional self-loop" + ) + +def reject_store(body, value, name): + pattern = rf"\bst\.global\.b32\s+\[[^\]]+\],\s*{value}\s*;" + if re.search(pattern, body): + raise SystemExit( + f"error: {name} retained unreachable post-drop store {value}" + ) + +diverging = ( + ("diverging_cfg_drop_probe", "LoopingDrop", 0x11111111, 0x22222222), + ("diverging_recursive_drop_probe", "RecursiveDrop", 0x33333333, 0x44444444), +) + +for name, type_name, before, after in diverging: + require_glue_call(kernel_body(name), type_name, name) + body = entry_body(name) + after_store = require_store(body, before, name) + reject_store(body, after, name) + require_nonreturning_loop(body, after_store, name) + print( + f"{name}: drop_glue::<{type_name}> call kept in drop_glue.ll, " + "pre-drop store followed by non-returning self-loop in drop_glue.ptx PASS" + ) + +# Precision control: an effect-free destructor whose branches join back +# together must still be proven a no-op, so no glue call may appear in the +# kernel and both of its stores survive. +name, before, after = ("reconverging_drop_probe", 0x55555555, 0x66666666) +reject_glue_call(kernel_body(name), name) +body = entry_body(name) +require_store(body, before, name) +require_store(body, after, name, label="post-drop") +print( + f"{name}: no drop glue call in drop_glue.ll, " + "pre-drop and post-drop stores kept in drop_glue.ptx PASS" +) + +print("drop_glue code shape: PASS") +VERIFY_PY 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 663fe0e1e4..818e5fb563 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/dst_metadata/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/Cargo.toml index 2ae94c410e..ff0a811628 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dst_metadata/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.lock index f1f4b2c026..d7a2d13df3 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.toml index ba275cffa1..989ff40cac 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/dynamic_smem/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } # Host-side utilities (CudaKernel trait, etc.) cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 f5ce40a552..af1b7f3494 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/elect_leader/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/Cargo.toml index a55c413dc1..6e5fa3c685 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/elect_leader/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 55fa067ced..cf9ff9bcf6 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/Cargo.toml index 9ef2b24ae3..a284c1e6a2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_abi_multi_payload/Cargo.toml @@ -23,4 +23,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 34e21c7c8e..bfd23bfe27 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/enum_array_match/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/Cargo.toml index 2478ad13b9..c86495b393 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_array_match/Cargo.toml @@ -22,4 +22,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 24abb1e836..f504b4ccae 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/enum_constant_provenance/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/Cargo.toml index a6716df8ad..5b300b2220 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_constant_provenance/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.lock index c073749d60..4504a40082 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.lock @@ -4,18 +4,18 @@ version = 4 [[package]] name = "aho-corasick" -version = "1.1.5" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" dependencies = [ "memchr", ] [[package]] name = "anyhow" -version = "1.0.104" +version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" [[package]] name = "bindgen" @@ -36,15 +36,15 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.119", + "syn", "which", ] [[package]] name = "bitflags" -version = "2.13.2" +version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" [[package]] name = "block-buffer" @@ -66,15 +66,15 @@ dependencies = [ [[package]] name = "cfg-if" -version = "1.0.5" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "clang-sys" -version = "1.9.1" +version = "1.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" dependencies = [ "glob", "libc", @@ -92,9 +92,9 @@ dependencies = [ [[package]] name = "crc32fast" -version = "1.5.2" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" dependencies = [ "cfg-if", ] @@ -128,19 +128,19 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", "prettyplease", "proc-macro2", "quote", - "syn 2.0.119", + "syn", ] [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,11 +151,11 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "syn", ] [[package]] @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,10 +185,11 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", - "syn 2.0.119", + "syn", ] [[package]] @@ -206,9 +208,9 @@ dependencies = [ [[package]] name = "either" -version = "1.19.0" +version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" [[package]] name = "enum_niche_payload" @@ -254,9 +256,9 @@ dependencies = [ [[package]] name = "glob" -version = "0.3.4" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" [[package]] name = "half" @@ -295,9 +297,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.14.2" +version = "2.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", "hashbrown 0.17.1", @@ -312,11 +314,17 @@ 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" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "lazycell" @@ -326,9 +334,9 @@ checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" [[package]] name = "libc" -version = "0.2.190" +version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce5d3ddc6d3fa000eb1536d85e147bfe31aacaba692ed6a876f95cb7c855be78" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libloading" @@ -357,9 +365,9 @@ checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" [[package]] name = "log" -version = "0.4.34" +version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" [[package]] name = "memchr" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -425,14 +442,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn 2.0.119", + "syn", ] [[package]] name = "proc-macro2" -version = "1.0.107" +version = "1.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" dependencies = [ "unicode-ident", ] @@ -443,18 +460,18 @@ version = "0.1.0" [[package]] name = "quote" -version = "1.0.47" +version = "1.0.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" dependencies = [ "proc-macro2", ] [[package]] name = "regex" -version = "1.13.1" +version = "1.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" dependencies = [ "aho-corasick", "memchr", @@ -464,9 +481,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.18" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" dependencies = [ "aho-corasick", "memchr", @@ -504,9 +521,9 @@ dependencies = [ [[package]] name = "serde" -version = "1.0.229" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" dependencies = [ "serde_core", "serde_derive", @@ -514,22 +531,35 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.229" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.229" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn 3.0.6", + "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]] @@ -551,20 +581,9 @@ checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" [[package]] name = "syn" -version = "2.0.119" +version = "2.0.117" 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" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" dependencies = [ "proc-macro2", "quote", @@ -573,22 +592,22 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.21" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.21" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn 3.0.6", + "syn", ] [[package]] @@ -599,9 +618,9 @@ checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" [[package]] name = "unicode-ident" -version = "1.0.26" +version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "version_check" @@ -711,20 +730,26 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "zerocopy" -version = "0.8.61" +version = "0.8.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "879fb705ce98c32e41ebdb970fbe1204f8492423b314c6ab0354c3e7b5542866" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.61" +version = "0.8.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "708882a28301d604fa039cc7727607a98d04c4b86dc76ec9cc683f805d709759" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "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/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.toml index fb6cca4f14..b843183e6f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } ptx-parse = { path = "../../../ptx-parse" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/src/main.rs index 627cf29383..269a0b5ec7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_niche_payload/src/main.rs @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA Corporation and AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ 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 1966983afe..df0cc37f42 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/enum_payload_addr/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/Cargo.toml index a121410b2c..ab9b766fa0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_payload_addr/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.lock index 4d1eb8f254..f7235d4900 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.toml index 4e99d7e931..55019aec40 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_state_transitions/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.lock index 795acb9cd4..a3895bfdff 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.toml index 2fb7ea0845..f06ad866bf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/enum_table_lookup/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 ae7af048ac..d1c009dd04 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.toml index 0ec1cebdd9..2b724a4917 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-macros = { path = "../../../../crates/cuda-macros" } cuda-device = { path = "../../../../crates/cuda-device" } cuda-host = { path = "../../../../crates/cuda-host" } 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 c7bd9ac758..5621f9b7d5 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_enum_bool_payload_addr/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_bool_payload_addr/Cargo.toml index bd5efceb2d..34e6811eaf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_bool_payload_addr/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_bool_payload_addr/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.lock index 5e1649b078..b6a117690d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.toml index bfa963cf25..b336af3def 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_pointer_overlap/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.lock index a8d8bd3db7..78af316b92 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.toml index 3ccdd08c61..83a612a458 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_enum_shared_pointer_layout/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_abi/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_generated_intrinsic_abi/Cargo.lock index 4d24441cc8..cf79f57d10 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -188,6 +189,7 @@ version = "99.0.2" name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -315,6 +317,12 @@ 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" @@ -421,6 +429,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -535,6 +552,19 @@ dependencies = [ "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" @@ -720,3 +750,9 @@ dependencies = [ "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/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 ac775b6804..f1e4b65bca 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -188,6 +189,7 @@ version = "0.1.0" name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -315,6 +317,12 @@ 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" @@ -421,6 +429,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -535,6 +552,19 @@ dependencies = [ "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" @@ -720,3 +750,9 @@ dependencies = [ "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/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 a471e93a9f..8f8b5bd07b 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -188,6 +189,7 @@ version = "0.1.0" name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -315,6 +317,12 @@ 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" @@ -421,6 +429,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -535,6 +552,19 @@ dependencies = [ "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" @@ -720,3 +750,9 @@ dependencies = [ "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/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 25462fd135..7c1967e280 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -188,6 +189,7 @@ version = "99.0.1" name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -315,6 +317,12 @@ 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" @@ -421,6 +429,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -535,6 +552,19 @@ dependencies = [ "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" @@ -720,3 +750,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.lock index 25827662ff..4b145cbb9b 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.toml index b92eed5fed..2e19b2df45 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_heap_alloc/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.lock index 5319a897c1..f13ccddf3d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.toml index 3de651b4af..b360107db0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_arch_intrinsic/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.lock index 360a02431e..896397563a 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.toml index effd001eec..7748c34fb4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_host_target_feature/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.lock index c04672d297..288baa5927 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.toml index 000882c68c..1cd2b85a3d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_param/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.lock new file mode 100644 index 0000000000..e526285811 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.lock @@ -0,0 +1,753 @@ +# 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 = "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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "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-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 = "error_kernel_shared_ptr_param" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[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 = [ + "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/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.toml new file mode 100644 index 0000000000..15fee6ddd1 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "error_kernel_shared_ptr_param" +version = "0.1.0" +edition = "2024" +authors = ["NVIDIA Corporation"] +license = "Apache-2.0" + +[workspace] + +[dependencies] +cuda-core = { path = "../../../../../cuda-core" } +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/README.md new file mode 100644 index 0000000000..0fcfdc5d95 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/README.md @@ -0,0 +1,32 @@ +# `error_kernel_shared_ptr_param` + +Negative test for a `#[kernel]` whose parameter is a struct that holds a +`SharedPtr`, a pointer that lowers to `addrspace(3)` by its type. + +A kernel receives its parameters in `.param` space, filled by the host at +launch. Shared memory is allocated per block by the device, so the host holds +no address in it to pass. `error_kernel_shared_param` pins the refusal of a +bare shared pointer parameter. This fixture pins the same refusal for a +pointer carried inside an aggregate: before the exporter looked through struct +fields and array elements, `Args { p: SharedPtr, n: u32 }` compiled, the +entry declared a 16-byte `.param`, and the kernel applied `cvta.shared.u64` to +the host-supplied bytes as if they were a shared address. + +The compiler must reject the signature at compile time instead: + +```bash +cargo oxide build error_kernel_shared_ptr_param +``` + +Expected diagnostic (pinned by `cuda-oxide/scripts/smoketest.sh`): + +```text +holds a pointer into shared memory +``` + +The diagnostic names the field path inside the parameter (`.0` here), which +keeps it distinct from the bare parameter's `is a pointer into shared memory`. +Structs of global (`addrspace(1)`) and constant (`addrspace(4)`) pointers stay +allowed as kernel parameters: the host allocates in those spaces, so it has an +address to supply. A `SharedPtr` belongs on the device side, created inside +the kernel from a `static mut SharedArray` or from `DynamicSharedArray`. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/src/main.rs new file mode 100644 index 0000000000..74f8b80c79 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_kernel_shared_ptr_param/src/main.rs @@ -0,0 +1,51 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Negative test: a `#[kernel]` parameter struct that holds a `SharedPtr`. +//! +//! A kernel receives its parameters in `.param` space, filled by the host at +//! launch. Shared memory is allocated per block by the device, so the host +//! holds no shared-memory address it could pass. A bare `SharedPtr` +//! parameter is refused for that reason (`error_kernel_shared_param` covers +//! the bare case with `*mut Barrier`), and a struct or array that carries one +//! must be refused the same way. Before the exporter looked inside +//! aggregates, `Args { p: SharedPtr, n: u32 }` compiled and the kernel +//! converted the host-supplied bytes with `cvta.shared.u64` as if they were a +//! shared address. +//! +//! Usage: +//! cargo oxide run error_kernel_shared_ptr_param +//! +//! Expected: the build FAILS with a diagnostic containing (pinned): +//! +//! ```text +//! holds a pointer into shared memory +//! ``` + +use cuda_device::{SharedPtr, kernel}; + +/// Kernel arguments that carry a typed shared pointer next to a plain field. +#[derive(Clone, Copy)] +pub struct Args { + pub p: SharedPtr, + pub n: u32, +} + +/// # Safety +/// +/// Never launched: the device compilation is expected to fail before a host +/// binary that could launch it exists. +#[kernel] +pub unsafe fn shared_ptr_param_kernel(args: Args) { + // BUG UNDER TEST: `args.p` points into shared memory, which the host + // cannot supply at launch. The refusal fires on the entry signature + // alone; the body deliberately does nothing else, so the diagnostic + // under test is the only error this fixture can produce. + let _ = args; +} + +fn main() { + println!("This negative example should fail during device compilation."); +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.lock index 91997eec33..73495fc325 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.toml index 41fc5521d8..9af52a9fd0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_missing_device_attr/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.lock index 2a2d387203..80b0290a9c 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.toml index 213d900319..5633559f80 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_set_discriminant_uninhabited/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.lock new file mode 100644 index 0000000000..b07deb9457 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.lock @@ -0,0 +1,753 @@ +# 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 = "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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "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-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 = "error_union_shared_pointer_addr" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[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 = [ + "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/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.toml new file mode 100644 index 0000000000..b9952904f9 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "error_union_shared_pointer_addr" +version = "0.1.0" +edition = "2024" +authors = ["NVIDIA Corporation"] +license = "Apache-2.0" + +[workspace] + +[dependencies] +cuda-core = { path = "../../../../../cuda-core" } +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/README.md new file mode 100644 index 0000000000..863bb7939b --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/README.md @@ -0,0 +1,35 @@ +# `error_union_shared_pointer_addr` + +Negative test for an in-place write to a union field that holds a +shared-memory pointer (`view.ptr = &raw mut STAGE`). + +A shared-memory pointer is 8 bytes wide on the PTX and legacy NVVM paths but +4 bytes wide on the modern NVVM path (`p3:32`). A union that holds one +therefore stores it as a generic pointer, which is 8 bytes on every path, and +the value paths (construct and extract) convert at that boundary. An in-place +write or a borrow of the field goes through the field's address instead. The +loads and stores made through that address carry the field's semantic type, +the raw shared pointer, so they would write or read a form of the pointer that +every other reader of the union treats as a generic address. On the modern +NVVM path that is a 4-byte store into an 8-byte slot. On the other paths the +shared-window offset comes back as a generic pointer, which is a wrong pointer +or an illegal address. + +The compiler refuses to form such an address. The same refusal covers +`&view.ptr`, `&mut view.ptr`, and writes into a struct nested in the field +(`view.inner.x = v`). Build the union by value and read its fields by value; +the `shared_ptr_transmute` example runs that path on the GPU. + +```bash +cargo oxide build error_union_shared_pointer_addr +``` + +Expected: the build FAILS, and the smoketest checks the log for this +diagnostic: + +```text +cannot hand out the in-place address of union field 0 of `SharedPtrBits` +``` + +together with the `canonical storage type` explanation from the +`mir.field_addr` lowering gate. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/src/main.rs new file mode 100644 index 0000000000..55265225bd --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_addr/src/main.rs @@ -0,0 +1,44 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Negative test: an in-place write to a union field that holds a +//! shared-memory pointer must fail closed. +//! +//! The union stores the pointer as a generic pointer so that its byte image is +//! the same on every target: a shared-memory pointer is 4 bytes wide on the +//! modern NVVM path and 8 bytes wide elsewhere. Construction and extraction by +//! value convert at that boundary. `view.ptr = p` instead stores through the +//! field's address with the field's own type, the raw shared pointer, into +//! bytes every other reader treats as a generic pointer. The compiler refuses +//! to form that address; `&view.ptr` and `&mut view.ptr` are refused the same +//! way. Build the union by value (`SharedPtrBits { ptr: p }`) and read its +//! fields by value instead. + +use cuda_device::{SharedArray, kernel}; + +static mut STAGE: SharedArray = SharedArray::UNINIT; + +#[allow(dead_code)] +union SharedPtrBits { + ptr: *mut SharedArray, + bits: u64, +} + +/// # Safety +/// +/// `out` must point to writable device memory for one `u64`, with no racing +/// access from another thread. +#[kernel] +pub unsafe fn write_shared_pointer_in_place(out: *mut u64) { + let mut view = SharedPtrBits { bits: 0 }; + view.ptr = &raw mut STAGE; + unsafe { + *out = view.bits; + } +} + +fn main() { + println!("This negative example should fail during device compilation."); +} 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 new file mode 100644 index 0000000000..e6c274596d --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/Cargo.lock @@ -0,0 +1,753 @@ +# 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 = "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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "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-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 = "error_union_shared_pointer_layout" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[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 = [ + "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/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/Cargo.toml new file mode 100644 index 0000000000..68d79940b5 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "error_union_shared_pointer_layout" +version = "0.1.0" +edition = "2024" +authors = ["NVIDIA Corporation"] +license = "Apache-2.0" + +[workspace] + +[dependencies] +cuda-core = { path = "../../../../../cuda-core" } +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/README.md new file mode 100644 index 0000000000..9cce8a79f7 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/README.md @@ -0,0 +1,31 @@ +# `error_union_shared_pointer_layout` + +Negative test for a union whose field is an array of seventeen +`*mut SharedArray<...>` values. + +Union fields that hold shared-memory pointers, directly, through structs, or +through bounded arrays, use target-stable generic-pointer storage, so the +union's byte image does not depend on whether the backend sizes a +shared-memory pointer at 8 bytes (PTX and legacy NVVM) or 4 bytes (modern +NVVM, `p3:32`). Array fields are rebuilt leaf by leaf at the union's construct +and extract boundary, with one extract, address-space cast and insert sequence +per pointer. To keep that expansion bounded, union lowering accepts at most 16 +array-expanded shared-pointer leaves per field, whether they come from one +array or from several arrays nested through structs. The enum payload path has +the same bound; see `error_enum_shared_pointer_layout`. + +This fixture exceeds the bound with 17 leaves. The compiler must reject it +instead of generating unbounded reconstruction or keeping the +target-dependent shared-pointer storage: + +```bash +cargo oxide build error_union_shared_pointer_layout +cargo oxide build error_union_shared_pointer_layout --emit-nvvm-ir --arch sm_90 +cargo oxide build error_union_shared_pointer_layout --emit-nvvm-ir --arch sm_100 +``` + +Expected diagnostic: + +```text +union field storage: arrays containing shared-memory pointers are not supported above the bounded rewrite limit; rewrite requires 17 pointer conversions, supported bound is 16 +``` diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/src/main.rs new file mode 100644 index 0000000000..7ada1170a7 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/error_union_shared_pointer_layout/src/main.rs @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Negative test: a union field holding an oversized array of shared-memory +//! pointers must fail closed when the rewrite to generic-pointer storage +//! would exceed the explicit code-shape bound. + +use cuda_device::{SharedArray, kernel}; + +static mut SHARED: SharedArray = SharedArray::UNINIT; + +#[allow(dead_code)] +union SharedPointerArrayBits { + pointers: [*mut SharedArray; 17], + words: [u64; 17], +} + +/// # Safety +/// +/// `out` must point to writable device memory for one `u64`, with no racing +/// access from another thread. +#[kernel] +pub unsafe fn shared_pointer_array_union(out: *mut u64) { + let shared: *mut SharedArray = &raw mut SHARED; + let view = SharedPointerArrayBits { + pointers: [shared; 17], + }; + unsafe { + *out = view.words[16]; + } +} + +fn main() { + println!("This negative example should fail during device compilation."); +} 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 2b169c6980..a305299bef 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/ex2_approx_f16/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/Cargo.toml index 82e021ff7b..7964fbadc2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ex2_approx_f16/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.lock index 87a392571e..7a093f77e0 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.toml index dd6a86a8c6..58ae125d45 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/exact_div_chunks/Cargo.toml @@ -22,4 +22,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 a5cf5dbb4b..d19a8f5d96 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/export_name_policy/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/Cargo.toml index d4ba585c94..ea79f961db 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/export_name_policy/Cargo.toml @@ -11,6 +11,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.lock index e45f61411e..f63c2d7694 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -316,6 +318,12 @@ 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" @@ -422,6 +430,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -536,6 +553,19 @@ dependencies = [ "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" @@ -721,3 +751,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.toml index 00fb7e3068..f5751e327a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_crate_closure/Cargo.toml @@ -13,4 +13,4 @@ members = ["helper-lib"] helper-lib = { path = "helper-lib" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 e35322740d..3ccac81def 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/extern_libdevice/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/Cargo.toml index 5b6b14297a..fadfa78f54 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/extern_libdevice/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 ba6233de9b..0ccebcd50f 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/f16_stress/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/Cargo.toml index f6596fee5c..27aeabbb47 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16_stress/Cargo.toml @@ -11,4 +11,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Primitive `f16` is nightly-only, so the shared crate gates its `DeviceCopy` # impl behind this feature; `half::f16` needs nothing extra. -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0", features = ["f16"] } +cuda-core = { path = "../../../../../cuda-core", features = ["f16"] } 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 d3dc0ec184..5a89dfb53e 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/f16x2_arith/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/Cargo.toml index 4436fea1dd..920b7e1db1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f16x2_arith/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.lock index 9df6fbcb48..1fef632193 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.toml index ce8ffe728f..8867e31dfb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/f32x2_arith/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.lock index 472983a662..5495cc6174 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.toml index d463d67756..f2a623a3a6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fcmp_ne_nan_smoke/Cargo.toml @@ -21,4 +21,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 de8cbf89e1..df607fcf01 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/field_array_assign/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/Cargo.toml index 026cb39121..7ae8748542 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/field_array_assign/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 82a53311d6..0b6f57b623 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/float_rounding/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/Cargo.toml index 8b6e992f87..bf7fee77ad 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/float_rounding/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 d9f5a36d7b..4458b7ec54 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/Cargo.toml index 8fe9fd3beb..6a3c7c0650 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fmaxmin_smoke/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 81d4b2aeb3..92d4321798 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/fn_ptr_naming/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/Cargo.toml index 6f3dad6650..e44115b467 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fn_ptr_naming/Cargo.toml @@ -21,4 +21,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 cc7a0b49d8..7507547040 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/Cargo.toml index 66ab77bf47..720b993bf1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/fp_special_literal_smoke/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 0aecde78e5..d21991bc5c 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/function_item_call/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/Cargo.toml index 019aeb159d..9fc5dc5011 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/Cargo.toml @@ -14,4 +14,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md index 62efd401c6..3a06d7c5af 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md @@ -7,3 +7,20 @@ Rust lowers `apply_once(plus_seven, x)` through a callable trait method such as function-item body, and MIR import must emit a direct call name that matches the collector/export naming policy. Otherwise the generated device IR references a callee symbol that was never emitted. + +The diverging-call regression also covers an ordinary direct `#[device]` +function returning `!`. Its body performs an observable store before looping +forever. A MIR `Call` with no normal successor must therefore keep the call and +end the caller block with `unreachable`; replacing the call itself with a trap +would erase the store. + +`verify-code-shape.sh` checks that contract in the pre-optimization LLVM IR, +before the middle-end can inline the diverging helper. It covers the three +call sites in `diverging_callable` that used to lower to a trap: two through +`apply_never` (a function item and a closure) and the direct call to +`write_then_diverge`. Each kept call must be followed by `unreachable`. When +the build also wrote PTX, the kernel entry must keep the sentinel store or a +call to the helper. The control is `panic_only`, a `-> !` function whose body +is nothing but a panic: the collector does not emit it, so its call site must +still lower to the device trap, in the IR and in the PTX. The divergent +branches are not launched; the host run exercises only the returning path. 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 68b442f644..f6bee22ce3 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 @@ -10,11 +10,16 @@ //! back to the concrete function body instead of emitting a dangling trait-shim //! callee symbol. The regression also covers `Fn`/`FnMut`, references, //! multiple and nested-tuple arguments, tuple returns, and type/const-generic -//! function items. Diverging function items and closures also verify that a -//! direct Rust call returning `!` ends its MIR block with `unreachable`. +//! function items. Diverging function items, closures, and an ordinary direct +//! `-> !` device helper also verify that a non-returning Rust call is emitted +//! before its MIR block ends with `unreachable`. `panic_only` is the control: +//! its `-> !` body is nothing but a panic, so the collector never emits it and +//! its call site must still lower to the device trap. //! `call_once_decoy` proves that an ordinary function whose name contains //! trait-like text is not mistaken for a callable-trait shim. +#![feature(never_type)] + use cuda_core::simt::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, device, kernel, thread}; @@ -57,6 +62,30 @@ fn spin_forever(_value: u32) -> ! { loop {} } +// Unlike `spin_forever`, this is reached by an ordinary direct MIR `Call`. +// The store makes replacing that call with an immediate trap observably wrong. +#[allow(clippy::empty_loop)] +#[inline(never)] +#[device] +unsafe fn write_then_diverge(ptr: *mut u32) -> ! { + unsafe { + ptr.write(0x1111_1111); + } + loop {} +} + +// Control for the retained trap. A `-> !` body that is nothing but a panic is +// a placeholder to the collector, so this function is never emitted and its +// call site must keep the device trap. `#[inline(never)]` keeps the call in +// the kernel's MIR, so the importer classifies this callee rather than the +// panic entry point inside it. The name must not contain a `panicking` module +// segment: that would trap the whole block before the call is classified. +#[inline(never)] +#[device] +fn panic_only() -> ! { + panic!("not collected") +} + #[device] fn apply_once u32>(f: F, x: u32) -> u32 { f(x) @@ -133,9 +162,10 @@ mod kernels { } } - /// Compile both diverging callable paths without entering them on-device. - /// A Rust call returning `!` has no successor block, so the importer must - /// emit the direct call followed by `mir.unreachable`. + /// Compile the diverging call paths without entering them on-device. + /// A Rust call returning `!` has no successor block. When the callee's + /// body is emitted, the importer must keep the direct call followed by + /// `mir.unreachable`; when it is not, the call site keeps the device trap. #[kernel] pub fn diverging_callable(flag: u32, mut out: DisjointSlice) { let idx = thread::index_1d(); @@ -148,6 +178,19 @@ mod kernels { #[allow(clippy::empty_loop)] apply_never(|_| loop {}, flag); } + if flag == 3 { + // The importer must preserve this ordinary direct call. The callee + // stores first and then diverges, so replacing it with a trap loses + // an observable side effect. + unsafe { + write_then_diverge(out.as_mut_ptr()); + } + } + if flag == 4 { + // Retained-trap control: the callee's body is only a panic, so the + // collector does not emit it and this call site must still trap. + panic_only(); + } if let Some(slot) = out.get_mut(idx) { *slot = 0xd1ce; } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/verify-code-shape.sh new file mode 100755 index 0000000000..a866ee8104 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/verify-code-shape.sh @@ -0,0 +1,104 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +llvm="${root}/function_item_call.ll" +ptx="${root}/function_item_call.ptx" + +test -s "${llvm}" + +extract_llvm_function() { + local needle="$1" + awk -v needle="${needle}" ' + /^define / && index($0, needle) != 0 { in_function = 1 } + in_function { print } + in_function && /^}/ { exit } + ' "${llvm}" +} + +kernel="$(extract_llvm_function '@diverging_callable(')" +test -n "${kernel}" + +# A kept diverging call must be followed by LLVM `unreachable` on the next +# line. A `--device-debug` build appends `, !dbg !N` to that line, so the +# match stops at the keyword. Monomorphized callees carry a mangled suffix +# after the Rust name (`apply_neverNvB2_12spin_foreverEB2_`), so the pattern +# allows symbol text on both sides of the name. +expect_kept_diverging_calls() { + local callee="$1" + local expected="$2" + local pattern="call void @[^ (]*${callee}[^ (]*\\(" + local found + found="$(grep -Ec "${pattern}" <<<"${kernel}" || true)" + if [[ "${found}" -ne "${expected}" ]]; then + echo "diverging_callable must keep ${expected} direct call(s) to ${callee}, found ${found}" >&2 + exit 1 + fi + local line next + while IFS= read -r line; do + next="$(sed -n "$((line + 1))p" <<<"${kernel}")" + if ! grep -Eq '^[[:space:]]*unreachable(,|[[:space:]]*$)' <<<"${next}"; then + echo "the direct diverging call to ${callee} at kernel line ${line} is not followed by LLVM unreachable" >&2 + exit 1 + fi + done < <(grep -nE "${pattern}" <<<"${kernel}" | cut -d: -f1) +} + +# The three call sites the importer used to replace with a trap: two through +# the generic `apply_never` (a function item and a closure) and the ordinary +# direct call to the store-then-loop helper. +expect_kept_diverging_calls 'apply_never' 2 +expect_kept_diverging_calls 'write_then_diverge' 1 + +# Retained-trap control. `panic_only` has a body that is nothing but a panic, +# which the collector does not emit, so its call site must still lower to the +# device trap and the callee must not be called. +trap_count="$(grep -Ec 'asm sideeffect "trap;"' <<<"${kernel}" || true)" +if [[ "${trap_count}" -ne 1 ]]; then + echo "diverging_callable must keep one trap for the panic-only callee, found ${trap_count}" >&2 + exit 1 +fi +if grep -Eq 'call void @[^ (]*panic_only[^ (]*\(' <<<"${kernel}"; then + echo "the panic-only callee must not be called: the collector does not emit its body" >&2 + exit 1 +fi + +helper="$(extract_llvm_function 'write_then_diverge')" +test -n "${helper}" +if ! grep -Eq 'store i32 286331153, ptr ' <<<"${helper}"; then + echo "write_then_diverge lost its observable sentinel store" >&2 + exit 1 +fi + +# The direct llc route also writes PTX; the NVVM route writes only `.ll` and +# `.target`. Check the PTX only when it exists and is at least as new as the +# IR, so a leftover from an older build is never checked. In the kernel entry +# the sentinel store must survive `opt` and `llc` (the helper is normally +# inlined) or the helper must still be called, and the panic-only control +# must keep its trap. +if [[ -s "${ptx}" && ! "${ptx}" -ot "${llvm}" ]]; then + extract_ptx_function() { + local pattern="$1" + awk -v pattern="${pattern}" ' + /^\.(visible )?\.(entry|func) / && $0 ~ pattern && $0 ~ /\($/ { in_function = 1 } + in_function { print } + in_function && /^}/ { exit } + ' "${ptx}" + } + entry="$(extract_ptx_function '[.]entry diverging_callable[(]')" + test -n "${entry}" + if ! grep -Eq 'st\.global\.b32[[:space:]]+\[[^]]+\], 286331153|call\.uni[[:space:]].*write_then_diverge' <<<"${entry}"; then + echo "the diverging_callable PTX entry neither stores the sentinel nor calls write_then_diverge" >&2 + exit 1 + fi + ptx_trap_count="$(grep -Ec '^[[:space:]]*trap;' <<<"${entry}" || true)" + if [[ "${ptx_trap_count}" -ne 1 ]]; then + echo "the diverging_callable PTX entry must keep one trap for the panic-only callee, found ${ptx_trap_count}" >&2 + exit 1 + fi +fi + +echo "function_item_call code shape: PASS" 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 6acb192412..767da36ce7 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/future_apis/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/Cargo.toml index 215421097d..6f57b6d7e4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-macros = { path = "../../../../crates/cuda-macros" } cuda-device = { path = "../../../../crates/cuda-device" } cuda-host = { path = "../../../../crates/cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md index 3b52278735..46757b69b6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/future_apis/README.md @@ -16,6 +16,7 @@ Tests new type-safe abstractions: `CuSimd` for SIMD values and `ManagedBar 1. **test_managed_barrier**: Single barrier typestate flow 2. **test_multi_barrier**: Multiple barriers (TMA + General) 3. **test_double_buffered_barriers**: Ping-pong pattern with two barriers +4. **test_managed_barrier_2d**: One barrier in a (32, 4, 1) block; every thread calls `init_by` and `inval_by`, and linear thread 127, that is thread (31, 3, 0), runs the instructions ## Key Concepts Demonstrated @@ -153,6 +154,10 @@ cargo oxide run future_apis All 32 threads completed 3 iterations ✓ Double-buffered PASSED +--- Test 6: ManagedBarrier in a (32, 4, 1) block --- + All 128 threads wrote 42 in a (32, 4, 1) block +✓ ManagedBarrier 2D PASSED + === ALL FUTURE APIS TESTS COMPLETED === ``` 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 d53c89ba93..55426c2f0a 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 @@ -11,11 +11,20 @@ //! //! Run: cargo oxide run future_apis +use cuda_device::cooperative_groups::{ThreadGroup, this_thread_block}; use cuda_device::cusimd::CuSimd; use cuda_device::{Barrier, GeneralBarrier, ManagedBarrier, TmaBarrier, Uninit}; use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; +/// Block of the 2D barrier test: 4 rows of 32 threads, one warp per row. +const BLOCK_2D_WIDTH: usize = 32; +const BLOCK_2D_HEIGHT: usize = 4; +const BLOCK_2D_THREADS: u32 = (BLOCK_2D_WIDTH * BLOCK_2D_HEIGHT) as u32; +/// Block-linear id of thread (31, 3, 0), the last thread of the block, so the +/// x and y terms of the id both count. No thread has `threadIdx.x` equal to it. +const BARRIER_2D_THREAD: u32 = BLOCK_2D_THREADS - 1; + // ============================================================================= // CUSIMD KERNELS // ============================================================================= @@ -231,6 +240,63 @@ mod kernels { *output_elem = sum; // Should be 3 } } + + /// Test kernel for `ManagedBarrier` in a 2D block of (32, 4, 1) threads. + /// + /// Every thread calls `init_by` and `inval_by`, as their contract asks: + /// both synchronize the whole block. The selector is the block-linear id + /// of thread (31, 3, 0), the last thread of the block, so the x and y + /// terms of the id both count. No thread has `threadIdx.x` equal to it, + /// so an election by `threadIdx.x` alone initializes nothing and the + /// launch fails with error 719. + /// + /// Thread 0 first fills the barrier word with all-ones. The store is a + /// determinism aid, not the detector: a missing init then faults on this + /// word instead of depending on what shared memory held at launch. + /// + /// Each thread writes 42 at its row-major 2D index when that index equals + /// its block rank, so the host also checks that `thread_rank()` numbers + /// the block x fastest, the same numbering the barrier election uses. + #[kernel] + pub fn test_managed_barrier_2d( + mut output: DisjointSlice>, + ) { + static mut BAR: Barrier = Barrier::UNINIT; + + let rank = this_thread_block().thread_rank() as usize; + + // Determinism aid, not the detector: see the kernel docs above + if rank == 0 { + // SAFETY: BAR is this block's own shared memory, one 8-byte aligned + // u64 that nothing in this launch has made an mbarrier object yet; + // the sync below orders the store before init_by. + unsafe { *(&raw mut BAR as *mut u64) = u64::MAX }; + } + thread::sync_threads(); + + // ALL threads create the Uninit handle (wraps the same shared static) + let bar = ManagedBarrier::::from_static(&raw mut BAR); + + // ALL threads call init_by - only linear thread 127 runs mbarrier.init + // init_by includes: fence + sync_threads; all threads get a Ready handle + let bar = unsafe { bar.init_by(BLOCK_2D_THREADS, BARRIER_2D_THREAD) }; + + // All threads arrive and wait + let token = bar.arrive(); + bar.wait(token); + + // ALL threads call inval_by - it synchronizes the block, then only + // linear thread 127 runs mbarrier.inval, then synchronizes again + let _dead = unsafe { bar.inval_by(BARRIER_2D_THREAD) }; + + // Write results - 42 when the row-major index is the block rank + if let Some(idx) = thread::index_2d::() { + let value = if idx.get() == rank { 42 } else { 1 }; + if let Some(output_elem) = output.get_mut(idx) { + *output_elem = value; + } + } + } } // ============================================================================= @@ -433,6 +499,42 @@ fn main() -> Result<(), Box> { } } + // ==================================================================== + // Test 6: ManagedBarrier in a 2D block + // ==================================================================== + println!("--- Test 6: ManagedBarrier in a ({BLOCK_2D_WIDTH}, {BLOCK_2D_HEIGHT}, 1) block ---"); + { + const N: usize = BLOCK_2D_WIDTH * BLOCK_2D_HEIGHT; + let mut output_dev = DeviceBuffer::::zeroed(&stream, N)?; + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (BLOCK_2D_WIDTH as u32, BLOCK_2D_HEIGHT as u32, 1), + shared_mem_bytes: 0, + }; + + // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + unsafe { module.test_managed_barrier_2d((stream).as_ref(), cfg, &mut output_dev) }?; + stream.synchronize()?; + + let output: Vec = output_dev.to_host_vec(&stream)?; + + // Every thread writes 42 once it passed the barrier and its block + // rank matched its row-major index. Any other value means the + // election this block layout exists to test went wrong, so it is an + // error, not a note. + if let Some(thread) = output.iter().position(|&v| v != 42) { + let got = output[thread]; + println!(" Results: {:?}", &output[..8.min(output.len())]); + return Err(format!( + "ManagedBarrier 2D: thread {thread} wrote {got} instead of 42 in a ({BLOCK_2D_WIDTH}, {BLOCK_2D_HEIGHT}, 1) block" + ) + .into()); + } + println!(" All {N} threads wrote 42 in a ({BLOCK_2D_WIDTH}, {BLOCK_2D_HEIGHT}, 1) block"); + println!("✓ ManagedBarrier 2D PASSED\n"); + } + println!("=== ALL FUTURE APIS TESTS COMPLETED ==="); Ok(()) } 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 056047ee7d..a16799a41f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/gemm/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.toml index c92ac5d159..2cf033aa72 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } # Host-side utilities (CudaKernel trait, etc.) cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 0aa6bf7395..e634080e85 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/gemm_sol/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/Cargo.toml index aab61a4f62..b8b19e4bc4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol/Cargo.toml @@ -15,5 +15,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } half = "2.4" 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 99a1c1c813..5704e7e7e2 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -313,6 +315,12 @@ 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" @@ -419,6 +427,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -533,6 +550,19 @@ dependencies = [ "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" @@ -718,3 +748,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.toml index 2d653f627c..b861d403da 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.toml @@ -14,6 +14,6 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } half = "2.4" ptx-parse = { path = "../../../ptx-parse" } 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..0749d8c40a 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 @@ -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 6b48d2b29b..d0a28e8c3a 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -312,6 +314,12 @@ 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" @@ -418,6 +426,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/gemm_views/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views/Cargo.toml index e165b7f14f..62aaa51b19 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } ptx-parse = { path = "../../../ptx-parse" } 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 e965b3570e..ab18e3fc3d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -188,6 +189,7 @@ version = "0.1.0" name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -316,6 +318,12 @@ 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" @@ -422,6 +430,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -536,6 +553,19 @@ dependencies = [ "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" @@ -721,3 +751,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/generated_intrinsics/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/Cargo.toml index a1576d1550..a21a0e435a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics/Cargo.toml @@ -12,4 +12,4 @@ license = "Apache-2.0" cuda-intrinsics = { path = "../../../cuda-intrinsics" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 65176bc3a3..adc7c6e4f4 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -188,6 +189,7 @@ version = "0.1.0" name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -316,6 +318,12 @@ 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" @@ -422,6 +430,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -536,6 +553,19 @@ dependencies = [ "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" @@ -721,3 +751,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/Cargo.toml index 3a402e3009..ea0e2e711c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_intrinsics_blackwell/Cargo.toml @@ -11,4 +11,4 @@ license = "Apache-2.0" cuda-intrinsics = { path = "../../../cuda-intrinsics" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 5454804be5..0d29f842d9 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 @@ -5,11 +5,12 @@ //! Coverage for generated high-target intrinsics. //! -//! The packed-FP8 and TF32 conversions are launched and their results checked -//! against the two formats, so the example reports what the device computed -//! rather than that it compiled. Both are bit-exact, which leaves no tolerance -//! to choose. The sparse MMA and TMA kernels stay compile-only: they need -//! operands and descriptors an example cannot supply. +//! The packed-FP8 conversions, from f32 pairs and from packed `bf16x2`, and +//! the TF32 conversions are launched and their results checked against the +//! formats, so the example reports what the device computed rather than that +//! it compiled. All are bit-exact, which leaves no tolerance to choose. The +//! sparse MMA and TMA kernels stay compile-only: they need operands and +//! descriptors an example cannot supply. //! //! Compile-only in CI, since it pins `sm_120a` and no runner has that device. //! Run it on one and the conversions are verified: @@ -22,10 +23,11 @@ use cuda_device::{ DisjointSlice, barrier::Barrier, convert::{ - cvt_rn_relu_satfinite_tf32_f32, cvt_rn_relu_tf32_f32, cvt_rn_satfinite_tf32_f32, - cvt_rn_tf32_f32, cvt_rna_satfinite_tf32_f32, cvt_rna_tf32_f32, - cvt_rz_relu_satfinite_tf32_f32, cvt_rz_relu_tf32_f32, cvt_rz_satfinite_tf32_f32, - cvt_rz_tf32_f32, + cvt_rn_relu_satfinite_tf32_f32, cvt_rn_relu_tf32_f32, cvt_rn_satfinite_e4m3x2_bf16x2, + cvt_rn_satfinite_e5m2x2_bf16x2, cvt_rn_satfinite_relu_e4m3x2_bf16x2, + cvt_rn_satfinite_relu_e5m2x2_bf16x2, cvt_rn_satfinite_tf32_f32, cvt_rn_tf32_f32, + cvt_rna_satfinite_tf32_f32, cvt_rna_tf32_f32, cvt_rz_relu_satfinite_tf32_f32, + cvt_rz_relu_tf32_f32, cvt_rz_satfinite_tf32_f32, cvt_rz_tf32_f32, }, cuda_module, kernel, thread, tma::{self, TmaDescriptor}, @@ -60,6 +62,28 @@ mod kernels { } } + /// Keeps every generated `bf16x2` to packed-FP8 conversion in device code. + /// + /// Launched, and its four results checked against the two formats. The + /// source is one register holding two bf16 values, low half first, and + /// the low half lands in the low byte of the result. + #[kernel] + pub fn compile_bf16x2_fp8_conversions(mut output: DisjointSlice, packed: u32) { + let values = [ + cvt_rn_satfinite_e4m3x2_bf16x2(packed), + cvt_rn_satfinite_relu_e4m3x2_bf16x2(packed), + cvt_rn_satfinite_e5m2x2_bf16x2(packed), + cvt_rn_satfinite_relu_e5m2x2_bf16x2(packed), + ]; + let start = thread::index_1d().get() * values.len(); + if start + values.len() <= output.len() { + for (offset, value) in values.into_iter().enumerate() { + // SAFETY: the bounds check covers this thread's unique slots. + unsafe { *output.get_unchecked_mut(start + offset) = value }; + } + } + } + /// Keeps every generated TF32 conversion in device code. /// /// Launched, and its ten results checked against the format. @@ -517,6 +541,53 @@ const FP8_CASES: [Fp8Case; 3] = [ }, ]; +/// One `bf16x2` case: the packed bf16 pair, low half in the low 16 bits, and +/// what each of the four conversions must return. +struct Bf16x2Case { + packed: u32, + e4m3: u16, + relu_e4m3: u16, + e5m2: u16, + relu_e5m2: u16, +} + +/// bf16 keeps the f32 exponent with a 7-bit mantissa, so the values of the +/// FP8 table are bf16-exact and the expected bytes carry over: 1.0 = 0x3f80, +/// 2.0 = 0x4000, -1.0 = 0xbf80, 448.0 = 0x43e0, 0.0 = 0x0000. The last two +/// cases add the bf16-only inputs: +inf (0x7f80) saturates to the largest +/// finite value, and a quiet NaN (0x7fc0) becomes 0x7f in both formats. +const BF16X2_CASES: [Bf16x2Case; 4] = [ + Bf16x2Case { + packed: 0x4000_3f80, + e4m3: 0x4038, + relu_e4m3: 0x4038, + e5m2: 0x403c, + relu_e5m2: 0x403c, + }, + // relu replaces the negative low lane with zero and leaves the high lane. + Bf16x2Case { + packed: 0x43e0_bf80, + e4m3: 0x7eb8, + relu_e4m3: 0x7e00, + e5m2: 0x5fbc, + relu_e5m2: 0x5f00, + }, + Bf16x2Case { + packed: 0x0000_7f80, + e4m3: 0x007e, + relu_e4m3: 0x007e, + e5m2: 0x007b, + relu_e5m2: 0x007b, + }, + Bf16x2Case { + packed: 0x3f80_7fc0, + e4m3: 0x387f, + relu_e4m3: 0x387f, + e5m2: 0x3c7f, + relu_e5m2: 0x3c7f, + }, +]; + /// TF32 keeps the f32 exponent and truncates the mantissa to 10 bits, so each /// value here is already representable and every rounding mode agrees. The /// four relu forms sit at indices 3, 5, 7 and 9 of the kernel's output. @@ -564,6 +635,27 @@ fn main() { } } + for case in &BF16X2_CASES { + let mut out = DeviceBuffer::::zeroed(&stream, 4).expect("alloc bf16x2 output"); + // SAFETY: the two arguments match the kernel's slice and scalar, and + // `out` holds the four results one thread writes. + unsafe { module.compile_bf16x2_fp8_conversions(&stream, cfg, &mut out, case.packed) } + .expect("bf16x2 launch failed"); + let got = out.to_host_vec(&stream).expect("bf16x2 readback failed"); + let want = [case.e4m3, case.relu_e4m3, case.e5m2, case.relu_e5m2]; + let names = ["e4m3x2", "relu_e4m3x2", "e5m2x2", "relu_e5m2x2"]; + for ((got, want), name) in got.iter().zip(want).zip(names) { + let mark = if *got == want { "ok" } else { "MISMATCH" }; + if *got != want { + failures += 1; + } + println!( + " bf16x2 {name:<12} (0x{:08x}) = 0x{got:04x} want 0x{want:04x} {mark}", + case.packed + ); + } + } + for (value, want_plain, want_relu) in TF32_CASES { let mut out = DeviceBuffer::::zeroed(&stream, 10).expect("alloc tf32 output"); // SAFETY: the two arguments match the kernel's slice and scalar, and 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 0e5a302952..2bd5d7e1ee 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -188,6 +189,7 @@ version = "0.1.0" name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -316,6 +318,12 @@ 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" @@ -422,6 +430,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -536,6 +553,19 @@ dependencies = [ "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" @@ -721,3 +751,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/generated_ldmatrix/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/Cargo.toml index fadb0977bb..d8c3e2b6ec 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generated_ldmatrix/Cargo.toml @@ -12,4 +12,4 @@ license = "Apache-2.0" cuda-intrinsics = { path = "../../../cuda-intrinsics" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 cf06cdd05d..ec5339be09 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/generic/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.toml index 4632c099ab..899b11b7f3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/generic/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 e44dd77b0c..751a0a2687 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/hashmap/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.toml index d23a0dbfea..e5f1070065 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 6a872c48e4..2646cae85d 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 @@ -159,7 +159,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -171,7 +171,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -182,7 +182,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -202,6 +202,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -215,6 +216,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -346,6 +348,12 @@ 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" @@ -452,6 +460,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -586,6 +603,19 @@ dependencies = [ "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" @@ -771,3 +801,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/hashmap_v2/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/Cargo.toml index e2ce14192c..53871da548 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v2/Cargo.toml @@ -28,6 +28,6 @@ path = "src/bin/bench.rs" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } hashbrown = "0.15" rayon = "1.10" 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 f5ae4facdd..b834c90e43 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 @@ -159,7 +159,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -171,7 +171,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -182,7 +182,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -202,6 +202,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -215,6 +216,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -346,6 +348,12 @@ 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" @@ -452,6 +460,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -586,6 +603,19 @@ dependencies = [ "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" @@ -771,3 +801,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/hashmap_v3/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/Cargo.toml index c4592846b0..a58577b00b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hashmap_v3/Cargo.toml @@ -27,6 +27,6 @@ path = "src/bin/bench.rs" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } hashbrown = "0.15" rayon = "1.10" 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 5279c6e7ef..1f7d33ddd1 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/hello_constant/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/Cargo.toml index e1528be7e1..fe6d668ae9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/hello_constant/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } # Host-side utilities (CudaKernel trait, etc.) cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 f6d933ccbc..9bd0891bc8 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/helper_fn/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/Cargo.toml index 840335717c..7cdd8d98b2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/helper_fn/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } # Host-side utilities (CudaKernel trait, etc.) cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 2add398677..dc0217c032 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 @@ -134,7 +134,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -172,7 +172,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -193,6 +193,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -206,6 +207,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -442,6 +444,12 @@ 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" @@ -557,6 +565,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.18", +] + [[package]] name = "parking_lot_core" version = "0.9.12" @@ -711,6 +728,19 @@ dependencies = [ "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" @@ -928,3 +958,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/host_closure/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/Cargo.toml index 8e534d4492..97a08b328d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure/Cargo.toml @@ -10,5 +10,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host", features = ["async"] } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } -cuda-async = { path = "../../../../../cuda-async", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } +cuda-async = { path = "../../../../../cuda-async" } 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 df77869652..cb58b6ab38 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/host_closure_target_feature/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/Cargo.toml index fe674badea..223e73cba4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/host_closure_target_feature/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 781249726e..879811c20f 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/iket_trace/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/Cargo.toml index aa5238658e..0aa1bd907d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/iket_trace/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.lock index 6cb6944b5c..00cd57c087 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 @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -152,7 +152,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -163,7 +163,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -183,6 +183,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -196,6 +197,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -344,6 +346,12 @@ 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" @@ -470,6 +478,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "png" version = "0.17.16" @@ -597,6 +614,19 @@ dependencies = [ "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" @@ -794,3 +824,9 @@ name = "zlib-rs" version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b268e58e7c693d7c271f93ffc4ba3b380412554231c85bf61ca7af91042a4112" + +[[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/examples/image_convolution/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.toml index ec74dd2bde..4782594b6b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/image_convolution/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } png = "0.17" 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 0686ce42dd..87f3c88f76 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/index2d_const/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/Cargo.toml index 22f866ad6c..32cd044d09 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/index2d_const/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.lock index 34f5c2907d..f6b0a95cd1 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.toml index 5f61206def..1525e91bbd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/index_field_assign/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.lock new file mode 100644 index 0000000000..5a05446079 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.lock @@ -0,0 +1,753 @@ +# 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 = "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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "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-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 = [ + "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 = "inline_intent" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[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/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.toml new file mode 100644 index 0000000000..00527d1967 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/Cargo.toml @@ -0,0 +1,25 @@ +cargo-features = ["profile-rustflags"] + +[package] +name = "inline_intent" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" + +# Mark as standalone crate (not part of parent workspace). +[workspace] + +[dependencies] +cuda-device = { path = "../../../cuda-device" } +cuda-host = { path = "../../../cuda-host" } +cuda-core = { path = "../../../../../cuda-core" } + +# Every helper here must reach LLVM as its own definition for its inline +# attribute to be observable. At default release settings rustc's MIR inliner +# folds the small `#[inline]` and `#[inline(always)]` helpers into the kernel +# first, and the shape check would have nothing to look at. Scoping the flag +# to this package keeps the shared dependency tree identical to every other +# example; a global RUSTFLAGS entry would re-key every dependency unit and +# force a full second dependency-tree build in CI. +[profile.release.package.inline_intent] +rustflags = ["-Zinline-mir=no"] 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 new file mode 100644 index 0000000000..037b086542 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/src/main.rs @@ -0,0 +1,135 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Inline intent reaches the GPU. +//! +//! Each `#[inline]` form on a helper reaches LLVM as the matching function +//! attribute: `#[inline]` as `inlinehint`, `#[inline(always)]` as +//! `alwaysinline` and `#[inline(never)]` as `noinline`. The one exception is +//! the `#[inline(never)]` that `#[device]` adds to a helper with its own +//! generics. That attribute only keeps each copy visible to the device +//! collector on the host side, so it is not carried to device code and +//! const-generic arguments keep folding into the kernel. +//! +//! The kernel computes the same arithmetic as the host reference through one +//! helper of each kind. `verify-code-shape.sh` checks the attribute on each +//! definition in the pre-opt LLVM IR and what opt made of it in the PTX: real +//! `call.uni` instructions for the `noinline` helpers, no function left for +//! the `alwaysinline` helper, and the const-generic factor folded to an +//! immediate. +//! +//! MIR inlining is off for this crate (see Cargo.toml) so that every helper +//! 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::{CudaContext, DeviceBuffer}; +use cuda_device::{DisjointSlice, cuda_module, device, kernel}; + +/// `#[inline(never)]` on a plain function: a real call boundary in the PTX. +#[inline(never)] +pub fn never_add(a: f32, b: f32) -> f32 { + a + b +} + +/// `#[inline]` on a plain function: an `inlinehint` for LLVM. +#[inline] +pub fn hint_mul(a: f32, b: f32) -> f32 { + a * b +} + +/// `#[inline(always)]` on a plain function: `alwaysinline`, so opt leaves no +/// function behind. A plain function is used because a non-generic `#[device]` +/// function is an exported symbol, and rustc ignores `#[inline(always)]` on +/// exported functions. +#[inline(always)] +pub fn always_sub(a: f32, b: f32) -> f32 { + a - b +} + +#[cuda_module] +mod kernels { + use super::*; + + /// `#[inline(never)]` on a non-generic `#[device]` function: also a real + /// call boundary. The macro adds no inline attribute of its own here. + #[inline(never)] + #[device] + pub fn never_dev(a: f32, b: f32) -> f32 { + a * 2.0 + b + } + + /// A const-generic `#[device]` function. The macro's own `#[inline(never)]` + /// is a host-side collector boundary and is not carried to device code, + /// so `N` folds into the kernel as an immediate. + #[device] + pub fn scale_by(x: f32) -> f32 { + x * (N as f32) + } + + #[kernel] + pub fn intent(input: &[f32], mut output: DisjointSlice) { + if let Some((slot, index)) = output.get_mut_indexed() { + *slot = reference(input, index.get()); + } + } +} + +/// The same arithmetic on host and device. `alwaysinline` folds it into the +/// kernel entry, so the entry itself carries the calls to the `noinline` +/// helpers. +#[inline(always)] +pub fn reference(input: &[f32], i: usize) -> f32 { + let x = input[i]; + let mut v = never_add(x, 1.0); + v = hint_mul(v, 2.0); + v = always_sub(v, 1.0); + v += kernels::never_dev(x, 1.0); + v + kernels::scale_by::<4>(x) +} + +fn main() { + const N: usize = 256; + let ctx = CudaContext::new(0).expect("Failed to create CUDA context"); + let stream = ctx.default_stream(); + let input: Vec = (0..N).map(|i| i as f32 * 0.25).collect(); + let input_dev = DeviceBuffer::from_host(&stream, &input).expect("Failed to upload input"); + let mut output_dev = + DeviceBuffer::::zeroed(&stream, N).expect("Failed to allocate output"); + + let module = kernels::load(&ctx).expect("Failed to load embedded CUDA module"); + // SAFETY: the launch covers N elements and both buffers hold N floats. + unsafe { + module.intent( + &stream, + LaunchConfig::for_num_elems(N as u32), + &input_dev, + &mut output_dev, + ) + } + .expect("Kernel launch failed"); + + let output = output_dev + .to_host_vec(&stream) + .expect("Failed to read output"); + let mismatches = output + .iter() + .enumerate() + .filter(|&(i, &got)| (got - reference(&input, i)).abs() > 1e-4) + .inspect(|&(i, &got)| { + eprintln!( + " mismatch at [{i}]: expected {}, got {got}", + reference(&input, i) + ); + }) + .count(); + println!("output[0..4] = {:?}", &output[0..4]); + if mismatches == 0 { + println!("SUCCESS: all {N} elements match the host reference"); + } else { + println!("FAILED: {mismatches} mismatches"); + std::process::exit(1); + } +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/verify-code-shape.sh new file mode 100755 index 0000000000..b7407d893d --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_intent/verify-code-shape.sh @@ -0,0 +1,134 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Each inline request must reach LLVM as its function attribute and leave the +# matching mark in the PTX. The pre-opt inline_intent.ll shows the attribute +# on each definition. The PTX shows what opt made of it: the `noinline` +# helpers stay real functions that the entry reaches with `call.uni`, the +# `alwaysinline` helper leaves no function behind, and the const-generic +# `#[device]` helper (whose macro-added `#[inline(never)]` is a host-side +# boundary only) folds its factor into an immediate. A regression in either +# direction fails compile-only CI. + +set -euo pipefail + +root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ll="${root}/inline_intent.ll" +ptx="${root}/inline_intent.ptx" + +test -s "${ll}" +test -s "${ptx}" + +# The one `define` line whose function name matches the pattern. +define_line() { + local pattern="$1" + local lines + lines="$(grep -E "^define .*@${pattern}\(" "${ll}" || true)" + if [[ -z "${lines}" ]]; then + echo "error: no definition matching ${pattern} in ${ll}" >&2 + exit 1 + fi + if [[ "$(wc -l <<<"${lines}")" -ne 1 ]]; then + echo "error: more than one definition matches ${pattern} in ${ll}:" >&2 + echo "${lines}" >&2 + exit 1 + fi + printf '%s\n' "${lines}" +} + +require_define_keyword() { + local description="$1" pattern="$2" keyword="$3" + local line + line="$(define_line "${pattern}")" + if ! grep -E "\) ${keyword} " <<<"${line}" >/dev/null; then + echo "error: ${description}: ${keyword} missing on: ${line}" >&2 + exit 1 + fi +} + +forbid_define_keywords() { + local description="$1" pattern="$2" + local line + line="$(define_line "${pattern}")" + if grep -E 'noinline|inlinehint|alwaysinline' <<<"${line}" >/dev/null; then + echo "error: ${description}: unexpected inline keyword on: ${line}" >&2 + exit 1 + fi +} + +# Print the body of one PTX entry. Waits for the header's `{` so a forward +# declaration (which ends in `;`) is skipped. +entry_body() { + local symbol="$1" + awk -v marker="${symbol}(" ' + !emit && !candidate && index($0, marker) && index($0, ".entry") { + candidate = 1 + } + candidate && $0 ~ /^[[:space:]]*;[[:space:]]*$/ { + candidate = 0 + next + } + candidate && $0 ~ /^[[:space:]]*\{[[:space:]]*$/ { + emit = 1 + candidate = 0 + } + emit { print } + emit && index($0, "End function") != 0 { exit } + ' "${ptx}" +} + +require_entry_shape() { + local symbol="$1" description="$2" pattern="$3" + if ! entry_body "${symbol}" | grep -E "${pattern}" >/dev/null; then + echo "error: missing ${description} in ${ptx}:${symbol}" >&2 + exit 1 + fi +} + +require_ptx_function() { + local description="$1" pattern="$2" + if ! grep -E "^[[:space:]]*(\.visible |\.weak )?\.func .*${pattern}\(" "${ptx}" >/dev/null; then + echo "error: missing ${description} in ${ptx}" >&2 + exit 1 + fi +} + +forbid_ptx_function() { + local description="$1" pattern="$2" + if grep -E "^[[:space:]]*(\.visible |\.weak )?\.func .*${pattern}\(" "${ptx}" >/dev/null; then + echo "error: ${description} is still a PTX function in ${ptx}" >&2 + exit 1 + fi +} + +# Pre-opt LLVM IR: one keyword per request, none on the generic #[device] +# instance whose #[inline(never)] belongs to the macro. +require_define_keyword "plain #[inline(never)] helper" 'inline_intent__never_add' noinline +require_define_keyword "non-generic #[inline(never)] #[device] helper" 'never_dev' noinline +require_define_keyword "plain #[inline] helper" 'inline_intent__hint_mul' inlinehint +require_define_keyword "plain #[inline(always)] helper" 'inline_intent__always_sub' alwaysinline +forbid_define_keywords "const-generic #[device] helper" 'scale_by[A-Za-z0-9_.$]*' + +# PTX: the noinline helpers survive opt as functions the entry calls. +require_ptx_function "plain noinline helper" 'inline_intent__never_add' +require_ptx_function "noinline #[device] helper" 'never_dev' +require_entry_shape intent "call to the plain noinline helper" 'inline_intent__never_add' +require_entry_shape intent "call to the noinline #[device] helper" 'never_dev' +if [[ "$(entry_body intent | grep -c 'call\.uni')" -lt 2 ]]; then + echo "error: the intent entry must call both noinline helpers with call.uni" >&2 + exit 1 +fi + +# PTX: alwaysinline leaves no function. The const-generic #[device] helper +# keeps a `.func` because every #[device] function stays exported through +# `@llvm.used`, but the entry must not call it: its factor 4.0 (0f40800000) +# is an immediate in the entry instead. +forbid_ptx_function "alwaysinline helper" 'inline_intent__always_sub' +if entry_body intent | grep -E 'scale_by' >/dev/null; then + echo "error: the intent entry still calls the const-generic #[device] helper in ${ptx}" >&2 + exit 1 +fi +require_entry_shape intent "folded const-generic factor" '0f40800000' + +echo "inline_intent code shape: PASS" 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 ebaeb8b7e0..9800de80bd 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/inline_ptx/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/Cargo.toml index 0dd7348737..b074c829e2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/inline_ptx/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.lock index ea07c279ad..38a2e99ccf 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.toml index 45240f8049..dc53282241 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/integer_minmax/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.lock index b9c688b9c5..87d0c79a58 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.toml index c58f305594..1be7dfaf94 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/interior_mutability/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.lock index 8746a26f4b..61cff5801b 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 @@ -117,7 +117,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -129,7 +129,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.toml index 3517c4d7d0..f2b57c59c9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/Cargo.toml @@ -21,5 +21,5 @@ bin = "scale-offset-device" resolver = "2" [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } sha2 = "0.10" 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 5b1e7cd88d..0671dfd50c 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -311,6 +313,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/interop_cubin_identity/device/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/device/Cargo.toml index b54b93c0bd..fcc03bf310 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/device/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/interop_cubin_identity/device/Cargo.toml @@ -18,4 +18,4 @@ resolver = "2" [dependencies] cuda-device = { path = "../../../../cuda-device" } cuda-host = { path = "../../../../cuda-host" } -cuda-core = { path = "../../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../../cuda-core" } 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 421a8e618e..15c94ec2a7 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 @@ -134,7 +134,7 @@ dependencies = [ [[package]] name = "cuda-async" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -149,7 +149,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -161,7 +161,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -172,7 +172,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -193,6 +193,7 @@ dependencies = [ "cuda-artifact-finalizer", "cuda-async", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -206,6 +207,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -432,6 +434,12 @@ dependencies = [ "either", ] +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + [[package]] name = "kernel_reference_validity" version = "0.1.0" @@ -556,6 +564,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.18", +] + [[package]] name = "parking_lot_core" version = "0.9.12" @@ -710,6 +727,19 @@ dependencies = [ "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" @@ -927,3 +957,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/kernel_reference_validity/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/Cargo.toml index c6b01baf4b..5f370c1c44 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/Cargo.toml @@ -1,3 +1,5 @@ +cargo-features = ["profile-rustflags"] + [package] name = "kernel_reference_validity" version = "0.1.0" @@ -8,6 +10,19 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host", features = ["async"] } + +# The `read_across_write` and `volatile_helper` kernels test that the reference +# facts on an entry parameter still reach loads inside a helper that LLVM +# inlines. At default release settings rustc's MIR inliner folds those small +# helpers into the kernel first, so LLVM would never see a call. `#[inline(never)]` +# is not the answer: since it is carried to LLVM as `noinline`, it would keep +# the helper a real PTX function and the loads would leave the entry. Switching +# MIR inlining off for this one crate keeps the calls for LLVM to inline. +# Scoping the flag to this package keeps the shared dependency tree identical +# to every other example; a global RUSTFLAGS entry would re-key every +# dependency unit and force a full second dependency-tree build in CI. +[profile.release.package.kernel_reference_validity] +rustflags = ["-Zinline-mir=no"] 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 32f62c9c19..5152441dd3 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 @@ -6,8 +6,8 @@ //! End-to-end LLVM-IR coverage for rustc-proven kernel reference validity. //! //! Build with `cargo oxide build kernel_reference_validity`, then run the host -//! binary (or `cargo oxide run kernel_reference_validity`) to inspect the -//! generated `.ll` parameter lines and execute valid slice-launch controls. +//! binary (or `cargo oxide run kernel_reference_validity`) for valid launches. +//! `verify-code-shape.sh` checks frontend attributes and optimized PTX loads. use cuda_device::{DisjointSlice, cuda_module, kernel, launch_contract, thread}; @@ -23,6 +23,19 @@ pub struct ByValue { pub pointer: *const f32, } +#[repr(C, align(16))] +pub struct GridConstants { + pub values: [u32; 4], +} + +/// # Safety +/// `output` must point to writable, aligned `u32` storage outside `constants`. +/// Concurrent invocations must not race on that storage. +#[kernel] +pub unsafe fn grid_constant(#[grid_constant] constants: &GridConstants, output: *mut u32) { + unsafe { *output = constants.values[0] }; +} + #[kernel] pub fn shared_ref(_value: &f32) {} @@ -41,6 +54,18 @@ pub fn align_one(_value: &u8) {} #[kernel] pub fn aligned_zst(_value: &AlignedZst) {} +// Interior mutability and pinning must follow rustc's flags, not `&` spelling. +#[kernel] +pub fn atomic_ref(_value: &core::sync::atomic::AtomicU32) {} + +pub struct PinnedValue { + pub value: u32, + pub pin: core::marker::PhantomPinned, +} + +#[kernel] +pub fn pinned_ref(_value: &mut PinnedValue) {} + #[kernel] pub fn raw_pointer(_value: *const f32) {} @@ -59,96 +84,63 @@ pub fn nvvm_slice(input: &[f32], mut output: DisjointSlice) { } } +/// With a valid shared input, the raw write cannot change the value read. +/// LLVM can reuse the first load once the reference guarantees reach it. +/// +/// # Safety +/// `scratch` and `output` must be writable, aligned `u32` pointers outside +/// `input`. Concurrent invocations must not race on their writes. #[kernel] -pub fn by_value(_value: ByValue) {} - -#[allow(improper_ctypes_definitions)] -#[kernel] -pub extern "C" fn c_reference(_value: &f32) {} - -fn kernel_header<'a>(llvm_ir: &'a str, name: &str) -> Result<&'a str, Box> { - let needle = format!("@{name}("); - llvm_ir - .lines() - .find(|line| line.trim_start().starts_with("define ") && line.contains(&needle)) - .ok_or_else(|| format!("missing LLVM kernel definition `{name}`").into()) -} - -fn require_reference_attrs( - llvm_ir: &str, - name: &str, - alignment: u64, -) -> Result<(), Box> { - let header = kernel_header(llvm_ir, name)?; - let expected = format!("nonnull align {alignment}"); - if !header.contains(&expected) { - return Err(format!( - "kernel `{name}` is missing `{expected}` on its reference parameter:\n{header}" - ) - .into()); +pub unsafe fn read_across_write(input: &u32, scratch: *mut u32, output: *mut u32) { + unsafe { + let before = read_input(input); + *scratch = 13; + *output = before.wrapping_add(read_input(input)); } - Ok(()) } -fn require_nonnull_without_alignment( - llvm_ir: &str, - name: &str, -) -> Result<(), Box> { - let header = kernel_header(llvm_ir, name)?; - if !header.contains("nonnull") || header.contains(" align ") { - return Err(format!( - "kernel `{name}` must carry nonnull without a redundant align-1 attribute:\n{header}" - ) - .into()); - } - Ok(()) +// MIR inlining is off for this crate (see Cargo.toml), so both calls reach +// LLVM, which sees the reference facts on the entry parameter when it inlines +// them. `#[inline(never)]` would keep this helper a real PTX function instead. +unsafe fn read_input(input: *const u32) -> u32 { + unsafe { *input } } -fn require_single_slice_pointer_fact( - llvm_ir: &str, - name: &str, -) -> Result<(), Box> { - let header = kernel_header(llvm_ir, name)?; - if header.matches("nonnull").count() != 1 || header.matches("align 4").count() != 1 { - return Err(format!( - "kernel `{name}` must annotate only the slice data pointer:\n{header}" - ) - .into()); - } - Ok(()) +/// Even an immutable input must be read twice when both reads are volatile. +/// Explicit `readonly` plus `noalias` made NVPTX lose `volatile` and allowed +/// ptxas to merge the reads. Withholding `readonly` preserves both accesses. +/// +/// # Safety +/// `output` must be writable, aligned, outside `input`, and owned by this thread. +#[kernel] +pub unsafe fn volatile_direct(input: &u32, output: *mut u32) { + let before = unsafe { core::ptr::read_volatile(input) }; + let after = unsafe { core::ptr::read_volatile(input) }; + unsafe { *output = before.wrapping_add(after) }; } -fn require_bare(llvm_ir: &str, name: &str) -> Result<(), Box> { - let header = kernel_header(llvm_ir, name)?; - if header.contains("nonnull") || header.contains(" align ") { - return Err(format!( - "kernel `{name}` unexpectedly acquired Rust-reference validity:\n{header}" - ) - .into()); - } - Ok(()) +/// The same guarantee must survive calls that LLVM later inlines. +/// +/// # Safety +/// `output` must be writable, aligned, outside `input`, and owned by this thread. +#[kernel] +pub unsafe fn volatile_helper(input: &u32, output: *mut u32) { + unsafe { *output = read_volatile_input(input).wrapping_add(read_volatile_input(input)) }; } -fn verify_generated_llvm_ir() -> Result<(), Box> { - let path = - std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("kernel_reference_validity.ll"); - let llvm_ir = std::fs::read_to_string(&path)?; - - require_reference_attrs(&llvm_ir, "shared_ref", 4)?; - require_reference_attrs(&llvm_ir, "unique_ref", 4)?; - require_single_slice_pointer_fact(&llvm_ir, "shared_slice")?; - require_single_slice_pointer_fact(&llvm_ir, "unique_slice")?; - require_single_slice_pointer_fact(&llvm_ir, "nvvm_slice")?; - require_nonnull_without_alignment(&llvm_ir, "align_one")?; - require_reference_attrs(&llvm_ir, "aligned_zst", 16)?; +// MIR inlining is off for this crate (see Cargo.toml), so the two calls reach +// LLVM and the loads land in the entry once LLVM inlines them. +fn read_volatile_input(input: &u32) -> u32 { + // SAFETY: the reference is aligned and initialized; reading does not mutate it. + unsafe { core::ptr::read_volatile(input) } +} - require_bare(&llvm_ir, "raw_pointer")?; - require_bare(&llvm_ir, "disjoint_slice")?; - require_bare(&llvm_ir, "by_value")?; - require_bare(&llvm_ir, "c_reference")?; +#[kernel] +pub fn by_value(_value: ByValue) {} - Ok(()) -} +#[allow(improper_ctypes_definitions)] +#[kernel] +pub extern "C" fn c_reference(_value: &f32) {} #[repr(C, align(32))] #[derive(Clone, Copy)] @@ -164,6 +156,16 @@ unsafe impl cuda_core::DeviceCopy for AlignedValue {} mod execution { use super::*; + #[kernel] + #[launch_contract(domain = 1, block = (32, 1, 1))] + pub fn add_inputs(a: &[f32], b: &[f32], mut output: DisjointSlice) { + let index = thread::index_1d(); + let i = index.get(); + if let Some(out) = output.get_mut(index) { + *out = a[i] + b[i]; + } + } + #[kernel] #[launch_contract(domain = 1, block = (1, 1, 1))] pub fn inspect_values(input: &[AlignedValue], mut output: DisjointSlice<[u64; 3]>) { @@ -197,6 +199,19 @@ fn verify_launch_packets() -> Result<(), Box> { let prepared = module.prepare_inspect_values(LaunchConfig1D::new(1, 1, 0))?; let prepared_zst = module.prepare_inspect_zsts(LaunchConfig1D::new(1, 1, 0))?; + // Two read-only parameters may share the same allocation across threads. + // Only the DisjointSlice output requires disjoint writable elements. + let values: Vec = (0..64).map(|i| i as f32).collect(); + let input = DeviceBuffer::from_host(&stream, &values)?; + let mut output = DeviceBuffer::::zeroed(&stream, values.len())?; + let add = module.prepare_add_inputs(LaunchConfig1D::new(2, 32, 0))?; + module.add_inputs(&stream, &add, &input, &input, &mut output)?; + assert_eq!( + output.to_host_vec(&stream)?, + values.iter().map(|value| value * 2.0).collect::>(), + "aliased read-only inputs with disjoint parallel output" + ); + for values in [vec![], vec![AlignedValue { values: [37; 8] }; 2]] { let input = DeviceBuffer::from_host(&stream, &values)?; let mut output = DeviceBuffer::<[u64; 3]>::zeroed(&stream, 1)?; @@ -261,10 +276,9 @@ fn verify_launch_packets() -> Result<(), Box> { } fn main() -> Result<(), Box> { - verify_generated_llvm_ir()?; verify_launch_packets()?; println!( - "SUCCESS: kernel reference validity checked in IR and sync/borrowed/owned async GPU launches" + "SUCCESS: kernel reference validity checked in sync/borrowed/owned async GPU launches" ); Ok(()) } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/verify-code-shape.sh new file mode 100644 index 0000000000..9763ca878d --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/kernel_reference_validity/verify-code-shape.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +set -euo pipefail + +root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# The frontend .ll precedes LLVM optimization; .opt.ll may infer extra facts. +python3 - "${1:-${root}/kernel_reference_validity.ll}" \ + "${2:-${root}/kernel_reference_validity.ptx}" <<'PY' +from pathlib import Path +import re +import sys + +ir = Path(sys.argv[1]).read_text() +attribute = re.compile(r"\b(?:nonnull|noalias|readonly)\b|\balign \d+|\bdereferenceable(?:_or_null)?\(\d+\)") + +def check(name, expected): + header = next((line for line in ir.splitlines() + if line.startswith("define ") and f"@{name}(" in line), None) + assert header is not None, f"missing kernel {name}" + # Match balanced parentheses: dereferenceable(N) is inside the arguments. + start = header.index(f"@{name}(") + len(name) + 2 + depth = 1 + end = start + while depth: + depth += (header[end] == "(") - (header[end] == ")") + end += 1 + params = header[start:end - 1].split(",") + actual = [set(attribute.findall(param)) for param in params] + assert actual == expected, f"{name}: expected {expected}, got {actual}\n{header}" + +reference = {"nonnull", "align 4", "noalias"} +bare = set() +check("shared_ref", [reference | {"dereferenceable(4)"}]) +check("unique_ref", [reference | {"dereferenceable(4)"}]) +check("shared_slice", [reference, bare]) +check("unique_slice", [reference, bare]) +check("nvvm_slice", [reference, bare, bare, bare]) +check("align_one", [{"nonnull", "noalias", "dereferenceable(1)"}]) +check("aligned_zst", [{"nonnull", "align 16", "noalias"}]) +check("atomic_ref", [{"nonnull", "align 4"}]) +check("pinned_ref", [{"nonnull", "align 4"}]) +check("grid_constant", [{"nonnull", "align 16", "noalias", "dereferenceable(16)"}, bare]) +check("read_across_write", [reference | {"dereferenceable(4)"}, bare, bare]) +for name in ("volatile_direct", "volatile_helper"): + check(name, [reference | {"dereferenceable(4)"}, bare]) +check("raw_pointer", [bare]) +check("disjoint_slice", [bare, bare]) +check("by_value", [bare]) +check("c_reference", [bare]) +print("PASS: exact frontend reference attributes and bare ABI controls") + +ptx = Path(sys.argv[2]).read_text() + +def global_loads(name): + entry = re.search(r"\.entry\s+" + re.escape(name) + r"\([^;]*?\n\{(.*?)^\}", + ptx, re.MULTILINE | re.DOTALL) + assert entry is not None, f"missing PTX entry {name}" + loads = re.findall(r"^\s*(?:@!?%[\w.$]+\s+)?(ld(?:\.[\w:]+)+)\s", + entry[1], re.MULTILINE) + return [load for load in loads if "global" in load.split(".")] + +# Runtime values alone cannot detect two volatile reads being merged into one. +# Check both the direct route and the route through an inlined helper. +for name in ("volatile_direct", "volatile_helper"): + loads = global_loads(name) + # LLVM and libNVVM may spell the same 32-bit integer load b32, u32, or s32. + assert len(loads) == 2 and all(re.fullmatch(r"ld\.volatile\.global\.[bus]32", load) + for load in loads), f"{name}: {loads}" + +# Keeping noalias still lets LLVM reuse an ordinary read across a disjoint write. +loads = global_loads("read_across_write") +assert len(loads) == 1 and re.fullmatch(r"ld\.global(?:\.nc)?\.[bus]32", loads[0]), \ + f"read_across_write: {loads}" +print("PASS: volatile reads stay separate; ordinary reads still combine") +PY 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 e0a335f508..940f65d37c 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ dependencies = [ "either", ] +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + [[package]] name = "lanemask_scan" version = "0.1.0" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/lanemask_scan/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/Cargo.toml index 0d45656dd8..e82094ff88 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/lanemask_scan/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 acedcabc8b..fd96ea2f4b 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/Cargo.toml index b51310232a..f7527e8b10 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_fadd/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 904c9c5d51..38e6fc926f 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/Cargo.toml index 9030008719..842304a725 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_atomic_rmw_cas/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 41da688d10..6a7f571741 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/Cargo.toml index 9f1a9259c4..dbdada1801 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/legacy_nvvm_pointer_shapes/Cargo.toml @@ -14,4 +14,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 285c0c6c88..4be2167e8e 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/libdevice_math/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/Cargo.toml index 84bfcd4fb2..21f2586113 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/libdevice_math/Cargo.toml @@ -15,4 +15,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 fe87b406fb..7838394dfe 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -424,6 +432,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -538,6 +555,19 @@ dependencies = [ "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" @@ -723,3 +753,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/libm_math/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/Cargo.toml index 4a21f4ed3e..408c4323a9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/libm_math/Cargo.toml @@ -19,5 +19,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } libm = "0.2" 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 175db3a78e..4e99d341bd 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/manual_launch_generic/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/Cargo.toml index caf96497e3..3c9e0cde0c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_generic/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 f402402704..4a2c19c285 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/Cargo.toml index c8b27ce73c..477d2c86ff 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/manual_launch_libdevice/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 8c33481722..236e1a127f 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/map_sum/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/Cargo.toml index f9be9aeb50..e01d2f87d6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/map_sum/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.lock index b60d260cf3..2e0d065824 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/math_atan/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.toml index f813df95c7..78e833a2ca 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_atan/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 bf4fc4cda4..51267d64d2 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/math_hyperbolic/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/Cargo.toml index 0fa502e77a..ed6c9725b2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_hyperbolic/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 9cbbb86a02..5c67e8bc57 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/math_no_mir_inline/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/Cargo.toml index 71511451d1..bc523e083b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_no_mir_inline/Cargo.toml @@ -12,7 +12,7 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } # `x.atan()` and every other float method is a one-line `#[inline]` wrapper in # `std`. At default release settings rustc's MIR inliner folds that wrapper 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 d32db3b7d9..c8bddba41d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/math_tan/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/Cargo.toml index 79c12e11af..113a3bc89e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/math_tan/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 8d14b54b1c..3b045b0c04 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 @@ -148,7 +148,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen 0.69.5", "libloading", @@ -160,7 +160,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -171,7 +171,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -191,6 +191,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -204,6 +205,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -331,6 +333,12 @@ 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" @@ -454,6 +462,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -574,6 +591,19 @@ dependencies = [ "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" @@ -759,3 +789,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/Cargo.toml index ff22aaa172..6f1c8ca902 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/Cargo.toml @@ -10,7 +10,7 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } # For runtime LTOIR generation (Option B) - optional libmathdx-sys = { path = "libmathdx-sys", optional = true } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/libmathdx-sys/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/libmathdx-sys/Cargo.toml index f51ad667db..a15ca2a78c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/libmathdx-sys/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mathdx_ffi_test/libmathdx-sys/Cargo.toml @@ -5,7 +5,7 @@ edition = "2024" description = "Raw FFI bindings to NVIDIA libmathdx (runtime code-generation for MathDx)" authors = ["Nihal Pasham "] license = "Apache-2.0" -repository = "https://github.com/NVlabs/cuda-oxide.git" +repository = "https://github.com/NVIDIA/cuda-rust" keywords = ["cuda", "ffi", "nvidia", "mathdx", "fft", "blas"] categories = ["external-ffi-bindings", "science"] 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 f8bd26cdd5..1c558ad854 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/mcast_barrier_test/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/Cargo.toml index 364fc8f2fb..35cde3d27d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mcast_barrier_test/Cargo.toml @@ -14,4 +14,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 9f0b7f6203..a126752583 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/mem_swap/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/Cargo.toml index 6f66fb3e79..4c4a34d945 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mem_swap/Cargo.toml @@ -19,4 +19,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 3fd8bac65e..af4602c219 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -424,6 +432,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -538,6 +555,19 @@ dependencies = [ "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" @@ -723,3 +753,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/Cargo.toml index 00cdb6a636..103837ab1c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mir_projected_call_destination/Cargo.toml @@ -13,5 +13,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } libm = "0.2" 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 1e680c9900..0070effb36 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/Cargo.toml index 8054015571..bea1f0dfb2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/mma_mxf8f6f4/Cargo.toml @@ -11,4 +11,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 978743e322..8d1dc68f4f 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/Cargo.toml index 4c521b448e..605cbeeb01 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ne_bytes_transmute/Cargo.toml @@ -21,4 +21,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 2be4cddb5c..d82c25246e 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/nested_index_assignment/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/Cargo.toml index 6fd9c4b586..db72405360 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nested_index_assignment/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.lock index 3de95b53ce..68d17a62ae 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.toml index 4ffdacf46d..71fec75276 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/numeric_stress/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 816249364b..58ef040c46 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/nvls_all_reduce/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/Cargo.toml index bbce4a6e42..72fef2884a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/nvls_all_reduce/Cargo.toml @@ -12,6 +12,6 @@ license = "Apache-2.0" # cargo oxide run nvls_all_reduce [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.lock index 74b48892c7..d6ccc1e244 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.toml index 32c9f4a75d..dc5a5bec75 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_transmute/Cargo.toml @@ -19,4 +19,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 8f05005c91..91de9b410d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/Cargo.toml index 2c86451a68..aa4dc14585 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_ref_unwrap_or/Cargo.toml @@ -20,4 +20,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 5dce21dac4..d164620488 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/option_unwrap_trap/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/Cargo.toml index b86ba0b26f..630d558e9d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/option_unwrap_trap/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 31ced58282..44d05c995d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -417,6 +425,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/ord_cmp/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/Cargo.toml index fef11f6c29..32ca865e01 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ord_cmp/Cargo.toml @@ -13,4 +13,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 09dddb0199..972c6cd402 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "packed_aggregate_abi" version = "0.1.0" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/Cargo.toml index bcf39f1dcf..46f948801c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_aggregate_abi/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 b3709e011e..af832aba00 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "packed_atomic_add" version = "0.1.0" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/packed_atomic_add/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/Cargo.toml index 6958fd95c0..dda30e5488 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/packed_atomic_add/Cargo.toml @@ -7,7 +7,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } half = "2.7" 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 393991e170..41a23e648d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "padded_constants" version = "0.1.0" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/padded_constants/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/Cargo.toml index d39a751bfe..74bc126f4c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/padded_constants/Cargo.toml @@ -25,4 +25,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 ba0ebf3c83..7e0a385ee1 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "partial_warp_reduce" version = "0.1.0" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/partial_warp_reduce/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/partial_warp_reduce/Cargo.toml index 1c0ca1b5a4..36ee8120ad 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/partial_warp_reduce/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/partial_warp_reduce/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.lock index f819cd7b3b..41c7b64525 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "pinned_overlap" version = "0.1.0" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.toml index a03e65a17e..991e7e775e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/Cargo.toml @@ -11,4 +11,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/README.md index e9ad3db610..88cad744e3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/README.md @@ -7,13 +7,13 @@ ways: 2. It compares a serialized pageable upload -> kernel -> download loop with a three-stream pipeline using rotating pinned host stagers and persistent device buffers. -Run it from the repository root: +From the cuda-oxide repository root: ```bash cargo oxide run pinned_overlap ``` -See the [book chapter](https://github.com/NVlabs/cuda-oxide/blob/main/cuda-oxide/cuda-oxide-book/async-programming/overlapping-transfers-and-compute.md) for the API and pipeline explanation. +See the [book chapter](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-oxide/cuda-oxide-book/async-programming/overlapping-transfers-and-compute.md) for the API and pipeline explanation. The program verifies every transformed element before printing its `SUCCESS` marker. It reports median CUDA-event bandwidth and wall-clock pipeline time so 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 78b2f77a2a..fdd244e8ef 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "place_read_fallbacks" version = "0.1.0" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/place_read_fallbacks/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/Cargo.toml index 5f1222767e..bcfeb6b965 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/place_read_fallbacks/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 713f0777ab..9cf6c64802 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "policy_config" version = "0.1.0" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/policy_config/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config/Cargo.toml index a5e05c012e..e0c38e0fa0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } ptx-parse = { path = "../../../ptx-parse" } 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 c84a6e15e3..84929b541d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/primitive_stress/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/Cargo.toml index 1ffe1aabfe..8cc0e7979e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/primitive_stress/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 a06ac1675b..9db04a93c7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/printf/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.toml index 214e223980..8434b1b225 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-macros = { path = "../../../../crates/cuda-macros" } cuda-device = { path = "../../../../crates/cuda-device" } cuda-host = { path = "../../../../crates/cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/README.md index 4466b14807..a985159e68 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/printf/README.md @@ -5,6 +5,9 @@ Tests the `gpu_printf!` macro for formatted output from GPU kernels. ## Features Tested - **Basic types**: integers (i32, u32, u64), floats (f32, f64), booleans +- **Default conversions**: `{}` prints its argument by type, `%llu` for a `u64` + above `i64::MAX`, `%f` for an `f32` or `f64` binding or slice element, `%p` + for a raw pointer; `{:.2}` on a float prints two decimals - **Format specifiers**: hex (`{:x}`, `{:X}`), octal (`{:o}`), scientific (`{:e}`, `{:E}`), compact (`{:g}`, `{:G}`) - **Flags**: left-justify (`{:-}`), sign (`{:+}`), space (`{: }`), alternate (`{:#}`), zero-pad (`{:0}`) - **Width and precision**: `{:8}`, `{:.2}`, `{:8.2}` @@ -26,7 +29,7 @@ cargo oxide run printf === Integer Tests === Signed i32: -42 Unsigned u32: 255 -Unsigned u64: 16045690984833335038 +Unsigned u64: 1311768467463790320 Hex (lower): ff Hex (upper): FF Hex with prefix: 0xff @@ -46,14 +49,30 @@ Compact (G): 1.23457E+06 ... (more tests) +--- Test 8: Default conversions --- +=== Default Conversion Tests === +u64::MAX: 18446744073709551615 +f32 binding: 3.141593 +f64 binding: 2.718282 +data[0]: 1.100000 +Precision .2 on f32: 3.14 +Zero-pad u64: 00000042 +Pointer: 0x7f3a9cc00000 +Mixed: -42 3.141593 18446744073709551615 2.72 ff + === ALL PRINTF TESTS PASSED === ``` +The pointer value changes from run to run. + ## How It Works The `gpu_printf!` macro: -1. **Compile-time**: Translates Rust format string (`{}`) to C format string (`%d`) +1. **Compile-time**: Translates the Rust format string (`{}`) to a C format + string. A placeholder without a type character takes its conversion from the + argument's type (`%d`, `%llu`, `%f`, `%p`); one with a type character + (`{:x}`, `{:.2f}`) casts the argument to the matching C type 2. **Compile-time**: Generates argument packing code with C vararg promotion rules 3. **Runtime**: Calls CUDA's `vprintf(format, args)` on the GPU 4. **Synchronization**: Output appears after `cudaDeviceSynchronize()` or similar 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 e6612da808..5900191f64 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 @@ -14,6 +14,8 @@ //! - Format specifiers: hex, octal, scientific, compact //! - Flags: left-justify, sign, alternate, zero-pad //! - Width and precision +//! - Default conversions chosen from the argument type: `{}` on a `u64` +//! above `i64::MAX`, on floats, on a slice element and on a raw pointer //! //! Run: cargo oxide run printf @@ -61,9 +63,9 @@ mod kernels { let small: f64 = 0.000123; gpu_printf!("=== Float Tests ===\n"); - gpu_printf!("f32 default: {:.6f}\n", pi); - gpu_printf!("f64 default: {:.6f}\n", e); - gpu_printf!("Precision .2: {:.2f}\n", pi); + gpu_printf!("f32 default: {}\n", pi); + gpu_printf!("f64 default: {}\n", e); + gpu_printf!("Precision .2: {:.2}\n", pi); gpu_printf!("Precision .6: {:.6f}\n", pi); gpu_printf!("Scientific (lower): {:e}\n", large); gpu_printf!("Scientific (upper): {:E}\n", large); @@ -165,6 +167,36 @@ mod kernels { gpu_printf!("false: {}\n", f); } } + + /// Test kernel: Conversions chosen from the argument type + /// + /// A placeholder without a type character prints its argument by type: + /// `%llu` for a `u64`, `%f` for a float, `%p` for a raw pointer. Each + /// line here once printed through `%lld`, so a `u64` above `i64::MAX` + /// came out negative and a float as a truncated integer. + #[kernel] + pub fn test_default_conversions(data: &[f32]) { + let tid = thread::index_1d().get(); + + if tid == 0 { + let max: u64 = u64::MAX; + let pi: f32 = 3.141_592_7; + let e: f64 = 2.718281828; + let small: u64 = 42; + let signed: i32 = -42; + let first: *const f32 = data.as_ptr(); + + gpu_printf!("=== Default Conversion Tests ===\n"); + gpu_printf!("u64::MAX: {}\n", max); + gpu_printf!("f32 binding: {}\n", pi); + gpu_printf!("f64 binding: {}\n", e); + gpu_printf!("data[{}]: {}\n", tid, data[tid]); + gpu_printf!("Precision .2 on f32: {:.2}\n", pi); + gpu_printf!("Zero-pad u64: {:08}\n", small); + gpu_printf!("Pointer: {}\n", first); + gpu_printf!("Mixed: {} {} {} {:.2} {:x}\n", signed, pi, max, e, 255u32); + } + } } // ============================================================================= @@ -255,6 +287,20 @@ fn main() -> Result<(), Box> { stream.synchronize()?; println!(); + // ==================================================================== + // Test 8: Conversions chosen from the argument type + // ==================================================================== + println!("--- Test 8: Default conversions ---"); + { + let data: Vec = vec![1.1, 2.2, 3.3, 4.4]; + let data_dev = DeviceBuffer::from_host(&stream, &data)?; + + // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + unsafe { module.test_default_conversions((stream).as_ref(), cfg, &data_dev) }?; + stream.synchronize()?; + } + println!(); + println!("=== ALL PRINTF TESTS PASSED ==="); Ok(()) } 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 88cde29662..17642ea7b6 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/projected_place_read/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/Cargo.toml index 1bcb8c2c53..3e8c6a0c47 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/projected_place_read/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 3ec04b9e64..46bf7cfb6b 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/proof_carrying_views/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views/Cargo.toml index 124d04e674..99ea121560 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views/Cargo.toml @@ -8,7 +8,7 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } ptx-parse = { path = "../../../ptx-parse" } 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 5b4640b69b..5b05e017c5 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/ptr_copy/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/Cargo.toml index b33fd99658..38435a50bc 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_copy/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.lock index c17c3d7a1c..3fa6d7132f 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.toml index 25a5e2a7c7..bee8a3e61a 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_offset_from/Cargo.toml @@ -19,4 +19,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 b9d72525d5..04e6a57cd5 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/ptr_without_provenance/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/Cargo.toml index 9af797f223..bf418c4686 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ptr_without_provenance/Cargo.toml @@ -16,4 +16,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 c6112cebe8..73188db112 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/reborrow/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.toml index c22aa0c126..d372977d04 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/reborrow/Cargo.toml @@ -11,7 +11,7 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } # The importer's `Rvalue::Reborrow` arm has a `Mutability::Mut` half that only # exists in unoptimized MIR: GVN folds a mutable reborrow into a plain copy at 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 fbe1e8cb36..d7a8c23a4e 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/redux_f32/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/Cargo.toml index de667cf7bd..f5d16d2aa0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_f32/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 ecbeeb54f8..7b05be8590 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/redux_minmax/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/Cargo.toml index 3f63ed4419..42563f6905 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_minmax/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 45927fe620..3cc9bb1335 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/redux_sum/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/Cargo.toml index ee960b4e78..efdcabe584 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/redux_sum/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 74c888841a..51d1860828 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/ref_index_projections/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/Cargo.toml index e11970280a..564544e3d1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_index_projections/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 94d0fa2102..c4a0a8bf84 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/ref_operand_mul/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/Cargo.toml index dca2fff37c..536b1688c1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/ref_operand_mul/Cargo.toml @@ -15,4 +15,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 159dcc98f3..ad83d723f6 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/repr_align_block_args/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/Cargo.toml index 2d34347cac..7705c7f9a0 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_align_block_args/Cargo.toml @@ -22,4 +22,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 10231a6158..07ce1da2d0 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -532,6 +549,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/repr_transparent_abi/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/Cargo.toml index 66bbb60bdf..76014df8c1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_transparent_abi/Cargo.toml @@ -10,5 +10,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } ptx-parse = { path = "../../../ptx-parse" } 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 c2f3be4056..79224818e2 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/Cargo.toml index f6f84f9b4f..44adc17356 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/README.md index 0d0caed544..3ef01ab567 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/repr_u32_enum_stride/README.md @@ -33,6 +33,7 @@ repr_c_enum [0, 1, 2, 3] PASS sparse_enum [0, 1000000, 1000000, 0] PASS neg_enum [-1, 0, -1, 0] PASS usize_enum [0, 1, 2, 3] PASS +align32_enum [0, 1, 1, 0] PASS RESULT: PASS - every enum shape strides and sign-extends correctly. ``` @@ -58,11 +59,17 @@ shape with one mechanism, which is why this example also tests: | `sparse_enum` | default repr, `B = 1_000_000` | u32 tag, stride 4 | | `neg_enum` | default repr, `N = -1` | SIGNED i8 tag, `sext` | | `usize_enum` | `#[repr(usize)]`, 4 unit variants | stride 8 | +| `align32_enum` | `#[repr(align(32))]`, 2 unit variants | u8 tag, stride 32, `align 32` on the load | `neg_enum` is the signedness half of the bug: the old model hardcoded an unsigned tag, so a memory-loaded `e as i32` zero-extended the `0xFF` byte to `255` instead of sign-extending it to `-1`. +`align32_enum` is the over-aligned shape: rustc pads the one-byte tag to 32 +bytes, and the LLVM storage type's zero-length alignment anchor is capped at +16 bytes, so both the stride and the `align 32` on the tag load must come +from rustc's layout rather than from the LLVM type. + Niched enums (for example `Option<&T>`) use a different rustc layout path: their otherwise-invalid payload value is recorded as the physical carrier, and cuda-oxide does not add a separate variant-count tag. 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 9964708571..d71f86e0c5 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 @@ -41,6 +41,13 @@ * memory-loaded `e as i32` must sign-extend to -1 * usize_enum `#[repr(usize)]` tag is pointer-width (8 bytes) * + * One more shape checks alignment rather than the tag: + * + * align32_enum `#[repr(align(32))]`, a u8 tag plus 31 padding bytes. The + * LLVM storage type's alignment anchor is capped at 16, so + * the 32-byte stride and the `align 32` on the tag load + * must both come from rustc's layout + * * Build: * cargo oxide run repr_u32_enum_stride */ @@ -116,6 +123,20 @@ pub enum UTag { // SAFETY: trivial repr(usize) POD, safe to copy device↔host. unsafe impl cuda_core::DeviceCopy for UTag {} +/// Fieldless `#[repr(align(32))]` enum: a u8 tag followed by 31 padding +/// bytes, so `*const Tag32` must stride by 32 over a byte buffer. The LLVM +/// storage type carries a zero-length alignment anchor capped at 16 bytes; +/// the 32 rides on the memory operations. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +#[repr(align(32))] +pub enum Tag32 { + A = 0, + B = 1, +} + +// SAFETY: fieldless POD enum, safe to copy device↔host. +unsafe impl cuda_core::DeviceCopy for Tag32 {} + const N: usize = 4; #[cuda_module] @@ -225,6 +246,23 @@ mod kernels { *slot = tag as u32; } } + + /// `#[repr(align(32))]` stride test: the enum is one tag byte plus 31 + /// padding bytes, so `*const Tag32` must stride by 32 over a byte buffer + /// and the tag load keeps its 32-byte alignment claim. + #[kernel] + pub fn read_via_align32_enum(input: &[u8], mut out: DisjointSlice) { + let idx = thread::index_1d(); + let i = idx.get(); + if i >= N { + return; + } + let base: *const Tag32 = input.as_ptr() as *const Tag32; + let tag = unsafe { *base.add(i) }; + if let Some(slot) = out.get_mut(idx) { + *slot = tag as u32; + } + } } /// Run one kernel over `input`, compare against `expected`, report PASS/FAIL. @@ -314,6 +352,24 @@ fn main() { unsafe { module.read_via_usize_enum(s, cfg, i, o) }.expect("launch") }); + // #[repr(align(32))]: one tag byte at the start of each 32-byte slot; + // stride 32 over a byte buffer. Device allocations are 256-byte aligned, + // so the 32-byte alignment the tag load claims holds. + let align32_tags: [u8; N] = [0, 1, 1, 0]; + let mut align32_in = vec![0u8; N * 32]; + for (slot, &tag) in align32_tags.iter().enumerate() { + align32_in[slot * 32] = tag; + } + let align32_expected: Vec = align32_tags.iter().map(|&tag| u32::from(tag)).collect(); + all_pass &= check( + "align32_enum", + &stream, + &align32_in, + &align32_expected, + // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + |s, cfg, i, o| unsafe { module.read_via_align32_enum(s, cfg, i, o) }.expect("launch"), + ); + println!(); if all_pass { println!("RESULT: PASS - every enum shape strides and sign-extends correctly."); 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 9cf1089479..1ed21433b6 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/row_width_slice/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/Cargo.toml index 32d3ccb88a..ef790e723f 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/row_width_slice/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 9d24a6b8d4..322099605b 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -306,6 +308,12 @@ 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" @@ -412,6 +420,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -536,6 +553,19 @@ dependencies = [ "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" @@ -721,3 +751,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/rustlantis-smoke/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/Cargo.toml index d7b780fec3..2b58008c05 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/rustlantis-smoke/Cargo.toml @@ -17,5 +17,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } fuzzer = { path = "../../../fuzzer" } 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 26d379b8e3..683cadeb40 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/Cargo.toml index 6b3f0a3192..61df35a3c8 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/scoped_atomic_load_store/Cargo.toml @@ -14,4 +14,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 a82c664e4e..8674c63f1c 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 @@ -54,10 +54,12 @@ mod kernels { /// `fence.acq_rel.sys` / `llvm.nvvm.membar.sys`, and the fused-fence /// SeqCst forms `fence.sc.gpu; st.release.gpu.b32` and `fence.sc.gpu; /// ld.acquire.gpu.b32`. + // These buffers are atomically mutated. Raw writable pointers avoid the + // immutability promise that &[u32] and &[u64] would give the compiler. #[kernel] pub fn publish_and_observe( - flags: &[u32], - wide: &[u64], + flags: *mut u32, + wide: *mut u64, mut observed: DisjointSlice, mut observed_wide: DisjointSlice, mut fence_ok: DisjointSlice, @@ -70,12 +72,10 @@ mod kernels { // SAFETY: `tid` and `next` are both < N, the buffers have N elements, // and every access to them in this kernel is atomic. - let mine = unsafe { DeviceAtomicU32::from_ptr((flags.as_ptr() as *mut u32).add(tid)) }; - let mine_wide = unsafe { DeviceAtomicU64::from_ptr((wide.as_ptr() as *mut u64).add(tid)) }; - let neighbour = - unsafe { DeviceAtomicU32::from_ptr((flags.as_ptr() as *mut u32).add(next)) }; - let neighbour_wide = - unsafe { DeviceAtomicU64::from_ptr((wide.as_ptr() as *mut u64).add(next)) }; + let mine = unsafe { DeviceAtomicU32::from_ptr(flags.add(tid)) }; + let mine_wide = unsafe { DeviceAtomicU64::from_ptr(wide.add(tid)) }; + let neighbour = unsafe { DeviceAtomicU32::from_ptr(flags.add(next)) }; + let neighbour_wide = unsafe { DeviceAtomicU64::from_ptr(wide.add(next)) }; // Publish. The release store is what an acquire reader synchronises with. mine.store(tid as u32 + 1, AtomicOrdering::Release); @@ -122,20 +122,27 @@ mod kernels { } } +// The unsafe launch owns writable-buffer access until stream completion. +fn atomic_storage(buffer: &mut DeviceBuffer) -> *mut T { + buffer.cu_deviceptr() as *mut T +} + fn main() { println!("=== Scoped atomic load/store and fences ===\n"); let ctx = CudaContext::new(0).expect("Failed to create CUDA context"); let stream = ctx.default_stream(); - let flags = DeviceBuffer::::zeroed(&stream, N).unwrap(); - let wide = DeviceBuffer::::zeroed(&stream, N).unwrap(); + let mut flags = DeviceBuffer::::zeroed(&stream, N).unwrap(); + let mut wide = DeviceBuffer::::zeroed(&stream, N).unwrap(); let mut observed = DeviceBuffer::::zeroed(&stream, N).unwrap(); let mut observed_wide = DeviceBuffer::::zeroed(&stream, N).unwrap(); let mut fence_ok = DeviceBuffer::::zeroed(&stream, N).unwrap(); let module = kernels::load(&ctx).expect("Failed to load embedded CUDA module"); - // SAFETY: one block of N threads, and every buffer has N elements. + // SAFETY: one block of N threads, every buffer has N elements, and buffers + // stay alive through completion. Flags/wide have no non-atomic accesses + // while the kernel runs; output copies synchronize the same stream. unsafe { module.publish_and_observe( &stream, @@ -144,8 +151,8 @@ fn main() { block_dim: (N as u32, 1, 1), shared_mem_bytes: 0, }, - &flags, - &wide, + atomic_storage(&mut flags), + atomic_storage(&mut wide), &mut observed, &mut observed_wide, &mut fence_ok, 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 b1cb5d2039..87763b5b63 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -531,6 +548,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/select_unpredictable/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/Cargo.toml index dab6abf19c..52cd150671 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/select_unpredictable/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.lock index a623091819..72726e3e91 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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 = "set_discriminant" version = "0.1.0" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.toml index aeb69db028..9406cbed42 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.lock index 158ecda0ac..5fbc5179e4 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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 = "set_discriminant_niche" version = "0.1.0" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.toml index e951a2c339..4944780439 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/set_discriminant_niche/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.lock index b790ed2ca1..c1225ffb45 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.toml index b1e0db543e..69aea86d4e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_debug/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 881232eacb..989f46f684 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -724,3 +754,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/shared_memory_helper/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/Cargo.toml index 9fdb9ebad7..7a0ead5e8d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/Cargo.toml @@ -18,5 +18,5 @@ heuristic-inline = ["shared-memory-helper-lib/heuristic-inline"] [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } shared-memory-helper-lib = { path = "helper-lib" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md index b9d5039ec9..6a2fbafa17 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md @@ -13,6 +13,58 @@ unsafe fn helper(smem: *mut u8, offset: usize) -> u64 { } ``` +The no-feature build also runs the same probe through `SharedPtr`, a +CTA-shared pointer typed as such (the feature builds fail MIR verification +before any probe runs): + +```rust +let smem = DynamicSharedArray::::shared_ptr(); +shared_ptr_probe(Region { base: smem, offset }); + +#[inline(never)] +unsafe fn shared_ptr_probe(region: Region) -> u64 { + let full = region.base.cast::().byte_add(1024).cast::().as_ptr(); + // ... +} +``` + +`SharedPtr` lowers to an `addrspace(3)` pointer by its type, so the +generic helper's parameter, nested in a struct, is one too, and the call +needs no address-space adaptation. + +Two more shapes run after that probe, in `shared_ptr_shapes_probe`: a struct +with two `SharedPtr` fields passed by value to the non-inlined +`write_pair`, which writes through both, and a `SharedPtr` returned from +the non-inlined `select_slot` and used as a store target. Each of 64 lanes +writes its own slots, then overwrites one slot of its partner lane through the +returned pointer, so the final values prove both stores landed. + +`verify-code-shape.sh` checks the LLVM IR and PTX the build leaves next to +this file: + +| Helper | Required code shape | +| --- | --- | +| `shared_ptr_probe` | first parameter is `ptr addrspace(3)`, and its call site passes one | +| `write_pair` | both `Pair` fields arrive as `ptr addrspace(3)` parameters | +| `select_slot` | return type is `ptr addrspace(3)` | +| `shared_ptr_probe` in PTX | a `.func` that the kernel calls once, with `mbarrier.init.shared` inside its body | +| whole module | no `cvta.to.shared` and no `isspacep` in the PTX | + +The PTX half reads the unoptimized build. `scripts/smoketest.sh` lists this +example in `NO_OPT_SHAPE_EXAMPLES`, so after the optimized build decides the +verdict it rebuilds once with `CUDA_OXIDE_NO_OPT=1` and runs the script on +those artifacts. To run the script by hand, build the same way first: + +```bash +CUDA_OXIDE_NO_OPT=1 cargo oxide build shared_memory_helper +bash crates/rustc-codegen-cuda/examples/shared_memory_helper/verify-code-shape.sh +``` + +Without the surviving `.func`, a module-wide search for `cvta.to.shared` +proves little: an optimized build inlines the helper, and the generic-pointer +control compiles to the same instructions. The call into the `.func` is where +a generic pointer would have needed a conversion. + The launch supplies 2048 dynamic shared bytes. Four independent CTAs each use one thread. Nine launches cover three data offsets and three input values, including wrapping arithmetic. Every CTA checks both a volatile shared-memory @@ -78,6 +130,11 @@ parameter has address space 0 (generic memory). The importer adapts that address-space difference before call verification. Pointee, mutability and pointer kind must still match. +With no features, the `SharedPtr` probe runs after the direct control and +reports `PASS shared-ptr: 9 launches, 36 CTA data/barrier checks`, and the +shapes probe reports `PASS shared-ptr-shapes: 3 launches, 384 slot checks`, +on the LLVM NVPTX route and with `--emit-nvvm-ir --arch sm_120`. + The dependency variant imports available MIR across crates; it does not test an opaque, separately linked device binary. The runtime checks one-thread barrier behavior, not synchronization between multiple threads or TMA traffic. 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 3abc8dc978..6b25119cb0 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 @@ -9,7 +9,8 @@ use cuda_core::simt::LaunchConfig; use cuda_core::{CudaContext, DeviceBuffer}; -use cuda_device::{DynamicSharedArray, cuda_module, kernel, launch_bounds, thread}; +use cuda_device::{DynamicSharedArray, SharedPtr, barrier::Barrier}; +use cuda_device::{cuda_module, kernel, launch_bounds, thread}; #[cfg(all(feature = "local-helper", feature = "cross-crate-helper"))] compile_error!("Select at most one helper feature; no features selects the direct control"); @@ -27,6 +28,82 @@ unsafe fn local_probe(smem: *mut u8, offset: usize) -> u64 { shared_memory_helper_lib::shared_probe_body!(smem, offset) } +/// The shared memory a typed probe works on, as a struct field. +#[derive(Clone, Copy)] +struct Region { + base: SharedPtr, + offset: usize, +} + +/// The probe body over a typed shared pointer, in a generic helper that is not +/// inlined. Unlike the raw `*mut u8` helpers above, it compiles: `SharedPtr` +/// is a CTA-shared pointer by its type, so the helper's parameter is one too, +/// and the barrier is addressed without converting a generic pointer. +/// +/// # Safety +/// As for `shared_probe_body!`: 2048 live shared bytes aligned to 1024, one +/// thread per CTA, and an initialized u64 at `offset`, below byte 1024. +#[inline(never)] +unsafe fn shared_ptr_probe(region: Region) -> u64 { + unsafe { + let bytes = region.base.cast::(); + let full: *mut Barrier = bytes.byte_add(1024).cast::().as_ptr(); + let data = bytes.byte_add(region.offset).cast::().as_ptr(); + cuda_device::barrier::mbarrier_init(full, 1); + let pending = !cuda_device::barrier::mbarrier_try_wait_parity(full, 0); + let _token = cuda_device::barrier::mbarrier_arrive(full); + let mut complete = 0u64; + let mut attempts = 0u32; + while attempts < 1024 { + if cuda_device::barrier::mbarrier_try_wait_parity(full, 0) { + complete = if pending { 1 } else { 2 }; + break; + } + attempts += 1; + } + let value = core::ptr::read_volatile(data); + core::ptr::write_volatile(data, value.wrapping_mul(3) ^ 0x1234_5678); + complete + } +} + +/// Two typed shared pointers in one struct, passed by value. +#[derive(Clone, Copy)] +struct Pair { + a: SharedPtr, + b: SharedPtr, +} + +/// Writes through both fields of a by-value `Pair`. The call boundary splits +/// the struct into two parameters; both must stay CTA-shared pointers. +/// +/// # Safety +/// Both fields must point at live shared memory with at least `index + 1` +/// `f32` slots. +#[inline(never)] +unsafe fn write_pair(pair: Pair, index: usize, value: f32) { + unsafe { + pair.a.add(index).as_ptr().write(value); + pair.b.add(index).as_ptr().write(value + 1000.0); + } +} + +/// Returns one slot of a `Pair` as a typed shared pointer. The return type +/// must stay a CTA-shared pointer so the caller's store needs no conversion. +/// +/// # Safety +/// As for `write_pair`. +#[inline(never)] +unsafe fn select_slot(pair: Pair, index: usize, second: bool) -> SharedPtr { + unsafe { + if second { + pair.b.add(index) + } else { + pair.a.add(index) + } + } +} + #[cuda_module] mod kernels { use super::*; @@ -52,6 +129,54 @@ mod kernels { *output.add(block * 2 + 1) = completed; } } + + /// The same probe through `SharedPtr`, a typed CTA-shared pointer, passed + /// to a non-inlined generic helper inside a struct. + #[kernel] + #[launch_bounds(1)] + pub fn shared_ptr_memory_probe(output: *mut u64, seed: u64, byte_offset: u64) { + let smem = DynamicSharedArray::::shared_ptr(); + let block = thread::blockIdx_x() as usize; + let offset = byte_offset as usize; + unsafe { + let data = smem.byte_add(offset).cast::().as_ptr(); + core::ptr::write_volatile(data, seed.wrapping_add(block as u64)); + let completed = shared_ptr_probe(Region { base: smem, offset }); + *output.add(block * 2) = core::ptr::read_volatile(data); + *output.add(block * 2 + 1) = completed; + } + } + + /// `SharedPtr` shapes a helper boundary must keep: a two-field struct of + /// typed shared pointers passed by value, and a typed shared pointer + /// returned from a helper and used as a store target. + #[kernel] + #[launch_bounds(64)] + pub fn shared_ptr_shapes_probe(output: *mut f32, seed: f32) { + const LANES: usize = 64; + let lane = thread::threadIdx_x() as usize; + let smem = DynamicSharedArray::::shared_ptr(); + // SAFETY: the launch supplies 2048 shared bytes, enough for the two + // 64-slot regions, and every lane writes only its own index and one + // slot of its partner lane, after a barrier. + unsafe { + let pair = Pair { + a: smem, + b: smem.add(LANES), + }; + let value = seed + lane as f32; + write_pair(pair, lane, value); + thread::sync_threads(); + // Each lane overwrites one slot of its partner lane through the + // pointer a helper returns; odd lanes pick the second field. + let partner = LANES - 1 - lane; + let slot = select_slot(pair, partner, lane % 2 == 1); + slot.as_ptr().write(-value); + thread::sync_threads(); + *output.add(lane) = *pair.a.add(lane).as_ptr(); + *output.add(LANES + lane) = *pair.b.add(lane).as_ptr(); + } + } } fn main() -> Result<(), Box> { @@ -79,41 +204,85 @@ fn main() -> Result<(), Box> { block_dim: (1, 1, 1), shared_mem_bytes: 2048, }; - let mut cases = 0; - for seed in [0u64, 17, u64::MAX] { - for offset in [0u64, 64, 248] { - // SAFETY: four independent CTAs, one thread each; output holds two - // u64s per CTA. Data offsets are aligned and disjoint from the - // barrier. Every launch receives a fresh shared-memory allocation. - unsafe { - module.shared_memory_probe( - &stream, - config, - output.cu_deviceptr() as *mut u64, - seed, - offset, - )?; - } - let got = output.to_host_vec(&stream)?; - for block in 0..BLOCKS { - let expected = seed.wrapping_add(block as u64).wrapping_mul(3) ^ 0x1234_5678; - assert_eq!( - got[block * 2], - expected, - "{variant}: data, seed={seed}, offset={offset}, block={block}" - ); - assert_eq!( - got[block * 2 + 1], - 1, - "{variant}: barrier, seed={seed}, offset={offset}, block={block}" - ); + for typed in [false, true] { + let variant = if typed { "shared-ptr" } else { variant }; + let mut cases = 0; + for seed in [0u64, 17, u64::MAX] { + for offset in [0u64, 64, 248] { + let out = output.cu_deviceptr() as *mut u64; + // SAFETY: four independent CTAs, one thread each; output holds + // two u64s per CTA. Data offsets are aligned and disjoint from + // the barrier. Every launch receives a fresh shared-memory + // allocation. + unsafe { + if typed { + module.shared_ptr_memory_probe(&stream, config, out, seed, offset)?; + } else { + module.shared_memory_probe(&stream, config, out, seed, offset)?; + } + } + let got = output.to_host_vec(&stream)?; + for block in 0..BLOCKS { + let expected = seed.wrapping_add(block as u64).wrapping_mul(3) ^ 0x1234_5678; + assert_eq!( + got[block * 2], + expected, + "{variant}: data, seed={seed}, offset={offset}, block={block}" + ); + assert_eq!( + got[block * 2 + 1], + 1, + "{variant}: barrier, seed={seed}, offset={offset}, block={block}" + ); + } + cases += 1; } - cases += 1; + } + println!( + "PASS {variant}: {cases} launches, {} CTA data/barrier checks", + cases * BLOCKS + ); + } + + const LANES: usize = 64; + let shapes_output = DeviceBuffer::::zeroed(&stream, LANES * 2)?; + let shapes_config = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (LANES as u32, 1, 1), + shared_mem_bytes: 2048, + }; + let seeds = [0.0f32, 17.0, 1024.0]; + for seed in seeds { + let out = shapes_output.cu_deviceptr() as *mut f32; + // SAFETY: one CTA of 64 lanes; output holds two f32 per lane, and the + // kernel uses 512 of the 2048 shared bytes supplied. + unsafe { module.shared_ptr_shapes_probe(&stream, shapes_config, out, seed)? }; + let got = shapes_output.to_host_vec(&stream)?; + for lane in 0..LANES { + let partner = LANES - 1 - lane; + // Lane `partner` overwrote one of this lane's slots: the second + // field when `partner` is odd, the first field otherwise. + let overwritten = -(seed + partner as f32); + let (want_a, want_b) = if partner % 2 == 1 { + (seed + lane as f32, overwritten) + } else { + (overwritten, seed + lane as f32 + 1000.0) + }; + assert_eq!( + got[lane], want_a, + "shared-ptr-shapes: field a, seed={seed}, lane={lane}" + ); + assert_eq!( + got[LANES + lane], + want_b, + "shared-ptr-shapes: field b, seed={seed}, lane={lane}" + ); } } println!( - "PASS {variant}: {cases} launches, {} CTA data/barrier checks", - cases * BLOCKS + "PASS shared-ptr-shapes: {} launches, {} slot checks", + seeds.len(), + seeds.len() * LANES * 2 ); Ok(()) } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/verify-code-shape.sh new file mode 100755 index 0000000000..5bc0dc0612 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_memory_helper/verify-code-shape.sh @@ -0,0 +1,103 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# `SharedPtr` keeps a CTA-shared pointer typed `addrspace(3)` across a call. +# +# LLVM IR: the non-inlined `shared_ptr_probe` helper takes one as its +# parameter and its call site passes one, the two fields of a by-value `Pair` +# reach `write_pair` as two such parameters, and `select_slot` returns one. +# +# PTX: `shared_ptr_probe` survives as a `.func` that the kernel calls once, +# and that `.func` initializes the barrier with `mbarrier.init.shared` on its +# parameter. No instruction converts a generic address to a shared one +# (`cvta.to.shared`) or asks which window an address is in (`isspacep`). The +# call into the `.func` is where a generic pointer would have needed one; a +# module-wide search alone proves little once the helper is inlined, because +# the generic-pointer control compiles to the same instructions. +# +# The PTX half reads the unoptimized build. scripts/smoketest.sh lists this +# example in NO_OPT_SHAPE_EXAMPLES and rebuilds it with CUDA_OXIDE_NO_OPT=1 +# before running this script; by hand, build the same way first. +set -euo pipefail + +dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ll="${dir}/shared_memory_helper.ll" +ptx="${dir}/shared_memory_helper.ptx" + +for file in "${ll}" "${ptx}"; do + [[ -f "${file}" ]] || { echo "FAIL: ${file} not found; build the example first" >&2; exit 1; } +done + +# --- LLVM IR ----------------------------------------------------------------- + +if ! grep -Eq '^define .*shared_ptr_probe.*\(ptr addrspace\(3\) ' "${ll}"; then + echo "FAIL: shared_ptr_probe does not take a ptr addrspace(3) parameter" >&2 + grep -E '^define .*shared_ptr_probe' "${ll}" >&2 || true + exit 1 +fi + +if ! grep -Eq 'call i64 @[^(]*shared_ptr_probe[^(]*\(ptr addrspace\(3\) ' "${ll}"; then + echo "FAIL: the call to shared_ptr_probe does not pass a ptr addrspace(3)" >&2 + grep -E 'call i64 @[^(]*shared_ptr_probe' "${ll}" >&2 || true + exit 1 +fi + +if ! grep -Eq '^define .*write_pair.*\(ptr addrspace\(3\) %[A-Za-z0-9_.]+, ptr addrspace\(3\) %[A-Za-z0-9_.]+, ' "${ll}"; then + echo "FAIL: write_pair does not take both Pair fields as ptr addrspace(3) parameters" >&2 + grep -E '^define .*write_pair' "${ll}" >&2 || true + exit 1 +fi + +if ! grep -Eq '^define ptr addrspace\(3\) @[^(]*select_slot' "${ll}"; then + echo "FAIL: select_slot does not return ptr addrspace(3)" >&2 + grep -E '^define .*select_slot' "${ll}" >&2 || true + exit 1 +fi + +# --- PTX --------------------------------------------------------------------- + +# The body of the PTX `.func` whose name contains the symbol: from the header +# line, past the parameter list, to the closing brace. A forward declaration +# ends in `;` before any `{` and is skipped. +func_body() { + awk -v needle="$1" ' + !emit && !candidate && index($0, ".func") && index($0, needle) { candidate = 1 } + candidate && $0 ~ /^[[:space:]]*;[[:space:]]*$/ { candidate = 0; next } + candidate && $0 ~ /^[[:space:]]*\{/ { emit = 1; candidate = 0 } + emit { print } + emit && $0 ~ /^\}/ { exit } + ' "${ptx}" +} + +probe_body="$(func_body shared_ptr_probe)" +if [[ -z "${probe_body}" ]]; then + echo "FAIL: no PTX .func for shared_ptr_probe in ${ptx}" >&2 + echo "hint: this half reads the unoptimized build; rebuild with" >&2 + echo " CUDA_OXIDE_NO_OPT=1 cargo oxide build shared_memory_helper" >&2 + exit 1 +fi + +if ! grep -Eq 'mbarrier\.init\.shared(::cta)?\.b64' <<<"${probe_body}"; then + echo "FAIL: the shared_ptr_probe .func does not initialize the barrier with mbarrier.init.shared" >&2 + exit 1 +fi + +# One call into the helper: the symbol on a `call` line, outside `.func` +# headers and comments. +calls="$(grep -E 'shared_ptr_probe' "${ptx}" | grep -vE '\.func|^[[:space:]]*//' | grep -cE '(^|[[:space:]])call' || true)" +if [[ "${calls}" -ne 1 ]]; then + echo "FAIL: expected one PTX call into shared_ptr_probe, found ${calls}" >&2 + grep -nE 'call.*shared_ptr_probe' "${ptx}" >&2 || true + exit 1 +fi + +for forbidden in 'cvta\.to\.shared' 'isspacep'; do + if grep -Eq "${forbidden}" "${ptx}"; then + echo "FAIL: PTX contains ${forbidden}: a generic address was converted or classified" >&2 + grep -nE "${forbidden}" "${ptx}" >&2 + exit 1 + fi +done + +echo "PASS: SharedPtr stays addrspace(3) across shared_ptr_probe, write_pair and select_slot; the shared_ptr_probe .func is called once and initializes its barrier in shared memory; no cvta.to.shared or isspacep" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock new file mode 100644 index 0000000000..35539d8e2b --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.lock @@ -0,0 +1,753 @@ +# 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 = "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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[package]] +name = "cuda-host" +version = "0.2.1" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "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-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 = [ + "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 = "shared_ptr_api" +version = "0.1.0" +dependencies = [ + "cuda-core", + "cuda-device", + "cuda-host", +] + +[[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/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.toml new file mode 100644 index 0000000000..f7fd6f703c --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "shared_ptr_api" +version = "0.1.0" +edition = "2024" +authors = ["NVIDIA Corporation"] +license = "Apache-2.0" +description = "SharedPtr through helper returns, references, closures, arrays, Option, shared storage and transmute" + +# Mark as standalone crate (not part of parent workspace) +[workspace] + +# Build with: +# cargo oxide run shared_ptr_api + +[dependencies] +# Device-side intrinsics (thread indexing, barriers, etc.) +cuda-device = { path = "../../../cuda-device" } + +# Host-side utilities (CudaKernel trait, etc.) +cuda-host = { path = "../../../cuda-host" } + +# Host-side CUDA runtime +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/README.md new file mode 100644 index 0000000000..52917d7c09 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/README.md @@ -0,0 +1,74 @@ +# shared_ptr_api + +`SharedPtr` is a raw pointer into the issuing block's shared memory whose +address space (`addrspace(3)`) is part of its type. A generic `*mut T` loses +that fact at the first helper boundary; a `SharedPtr` keeps it through +parameters, return values, struct fields, closures and arrays, so the compiler +can use the shared-memory forms of the instructions that take it. This example +moves one through each of those boundaries and checks every value on the host. + +| Kernel | Boundary | Check | +| :------------------ | :-------------------------------------------------------------- | :-------------------------------------- | +| `static_round_trip` | `SharedArray::shared_ptr`, `from_raw(as_ptr())`, negative `offset` | each thread reads back its own slot | +| `closure_capture` | a closure that captured a `SharedPtr`, called through a generic `#[inline(never)]` helper | `3 * tid` | +| `addr_u32_deltas` | `addr_u32` of `shared_ptr_at(64)`, `add(3)`, `byte_add(7)`, and of `get_raw()` through `cvta_generic_to_shared_u32` | `[64, 12, 1, 7]` | +| `behind_references` | `&Region` (a struct holding the pointer) and `&mut SharedPtr` through non-inlined helpers | three reads summed | +| `returned_pointer` | a non-inlined helper returning `SharedPtr` | odd lanes read the pool, even the array | +| `array_by_value` | `[SharedPtr; 2]` passed by value | `7 + tid` | +| `struct_pointee` | `SharedPtr` and `cast::()` to its first field | `a + b + a` | +| `byte_views` | `byte_add` into a `u8` pool, viewed as `f64` and `u16` | `0.25 * tid + tid + 7` | +| `option_pointer` | `Option>` returned and passed | `Some` stores, `None` does not | +| `stored_in_shared` | `SharedArray, 64>`: pointers kept in shared memory | the mirrored slot's value | +| `transmuted` | `transmute` to `usize` and back in non-inlined helpers | the store lands and the stride is 4 | + +The helpers carry `#[inline(never)]`, so each boundary is a real call. In the +LLVM IR the default build leaves next to this file, `smem_fill`, `smem_load`, +`pick`, `maybe` and the closure's `apply` take `ptr addrspace(3)`, +`store_via_array` takes `[2 x ptr addrspace(3)]`, `pick` and `from_bits` +return `ptr addrspace(3)`, and `SLOTS` is a `[64 x ptr addrspace(3)]` global +in address space 3. The PTX converts no generic address to a shared one +except in `from_bits`: the integer view of a shared pointer is its generic +address, as for unions, so the transmute back is one `cvta.to.shared`. + +Structs that mix a `SharedPtr` with a narrower field, such as +`{ p: SharedPtr, x: u32 }`, stay out of this example. On the libNVVM +route a shared pointer is 4 bytes wide, and such a struct is not yet +supported in arrays or in shared memory; the book's shared memory chapter +lists this limit. + +## Build and run + +```bash +cargo oxide run shared_ptr_api +``` + +The compile-only lane of `scripts/smoketest.sh` also compiles this example +through libNVVM (`NVVM_VERIFY_EXAMPLES`), where the shared pointer is 4 bytes +wide: + +```bash +cargo oxide emit-ltoir shared_ptr_api --arch=sm_120 +``` + +## Expected output + +```text +=== SharedPtr API (64 threads) === + ok static_round_trip + ok closure_capture + ok addr_u32_deltas + ok behind_references + ok returned_pointer + ok array_by_value + ok struct_pointee + ok byte_views + ok option_pointer + ok stored_in_shared + ok transmuted +SUCCESS: shared_ptr_api, 11 kernels, 0 mismatches +``` + +## Hardware requirements + +Any CUDA GPU; the kernels use plain shared memory and `sync_threads()`. The +launch supplies 1024 bytes of dynamic shared memory. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/src/main.rs new file mode 100644 index 0000000000..12ba608b2b --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_api/src/main.rs @@ -0,0 +1,535 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! `SharedPtr` across every boundary a helper call can put it behind. +//! +//! `SharedPtr` is a raw pointer into the issuing block's shared memory +//! whose address space is part of its type. Each kernel here moves one +//! through a boundary where a generic `*mut T` would lose that fact, or where +//! the lowering has to rebuild the value from parts: +//! +//! - a `from_raw` round trip and a negative `offset` on a static array, +//! - a closure that captured one, called through a generic helper, +//! - `addr_u32` deltas against `cvta_generic_to_shared_u32(get_raw())`, +//! - `&Region` and `&mut SharedPtr` through non-inlined helpers, +//! - a non-inlined helper that returns one, +//! - `[SharedPtr; 2]` passed by value, +//! - a struct pointee and a `cast` to its first field, +//! - `f64` and `u16` views of a dynamic byte pool, +//! - `Option>` returned and passed, +//! - `SharedPtr` values stored in shared memory itself, +//! - `transmute` to `usize` and back. +//! +//! The host checks every value and prints one SUCCESS line. Structs that mix +//! a `SharedPtr` with a narrower field stay out of this example: on the +//! libNVVM route a shared pointer is 4 bytes wide, and such a struct is not +//! yet supported in arrays or in shared memory. +//! +//! Usage: +//! cargo oxide run shared_ptr_api + +use cuda_core::simt::LaunchConfig; +use cuda_core::{CudaContext, DeviceBuffer}; +use cuda_device::shared::cvta_generic_to_shared_u32; +use cuda_device::{DynamicSharedArray, SharedArray, SharedPtr, cuda_module, kernel, thread}; + +const THREADS: usize = 64; + +/// A pointee with two fields; a `cast` to `f32` reaches the first. +#[derive(Clone, Copy)] +#[repr(C)] +struct Pair { + a: f32, + b: f32, +} + +/// A typed shared pointer next to an element offset, read through `&Region`. +#[derive(Clone, Copy)] +struct Region { + base: SharedPtr, + offset: usize, +} + +static mut TILE_F32: SharedArray = SharedArray::UNINIT; +static mut TILE_U32: SharedArray = SharedArray::UNINIT; +static mut PAIRS: SharedArray = SharedArray::UNINIT; +static mut SLOTS: SharedArray, THREADS> = SharedArray::UNINIT; + +/// Stores through a typed shared pointer inside a non-inlined generic helper. +/// +/// # Safety +/// `p.add(i)` must be in bounds of live shared memory. +#[inline(never)] +unsafe fn smem_fill(p: SharedPtr, i: usize, v: T) { + unsafe { p.add(i).as_ptr().write(v) } +} + +/// Loads through a typed shared pointer inside a non-inlined generic helper. +/// +/// # Safety +/// `p.add(i)` must be in bounds of initialized shared memory. +#[inline(never)] +unsafe fn smem_load(p: SharedPtr, i: usize) -> T { + unsafe { p.add(i).as_ptr().read() } +} + +/// Returns one of two typed shared pointers. The return type lowers to +/// `ptr addrspace(3)`. +#[inline(never)] +fn pick(a: SharedPtr, b: SharedPtr, second: bool) -> SharedPtr { + if second { b } else { a } +} + +/// Calls a closure that captured a `SharedPtr` by value. +#[inline(never)] +fn apply(f: F, i: usize) { + f(i) +} + +/// Reads through `&Region`: a deref, a field of `Region`, then the pointer. +/// +/// # Safety +/// `r.base.add(r.offset + i)` must be in bounds of initialized shared memory. +#[inline(never)] +unsafe fn read_region(r: &Region, i: usize) -> T { + unsafe { r.base.add(r.offset + i).as_ptr().read() } +} + +/// Advances a `SharedPtr` in place through `&mut SharedPtr`. +/// +/// # Safety +/// The advanced pointer must stay inside the same allocation. +#[inline(never)] +unsafe fn bump(slot: &mut SharedPtr, by: usize) { + unsafe { *slot = slot.add(by) } +} + +/// Stores through one of two `SharedPtr`s passed in an array by value. +/// +/// # Safety +/// `ptrs[which].add(i)` must be in bounds of live shared memory. +#[inline(never)] +unsafe fn store_via_array(ptrs: [SharedPtr; 2], which: usize, i: usize, v: T) { + unsafe { ptrs[which].add(i).as_ptr().write(v) } +} + +/// Stores through an `Option>` and reports whether it did. +/// +/// # Safety +/// A `Some` pointer must be in bounds of live shared memory at `i`. +#[inline(never)] +unsafe fn store_opt(p: Option>, i: usize, v: u32) -> u32 { + match p { + Some(p) => { + unsafe { p.add(i).as_ptr().write(v) }; + 1 + } + None => 0, + } +} + +/// Wraps a `SharedPtr` in `Some` or drops it, in a non-inlined helper. +#[inline(never)] +fn maybe(p: SharedPtr, take: bool) -> Option> { + if take { Some(p) } else { None } +} + +/// A `SharedPtr` as its address bits. +#[inline(never)] +fn to_bits(p: SharedPtr) -> usize { + // SAFETY: `SharedPtr` is `repr(transparent)` over a thin pointer. + unsafe { core::mem::transmute::, usize>(p) } +} + +/// Address bits back to a `SharedPtr`. +/// +/// # Safety +/// `bits` must come from `to_bits` on a live `SharedPtr`. +#[inline(never)] +unsafe fn from_bits(bits: usize) -> SharedPtr { + unsafe { core::mem::transmute::>(bits) } +} + +#[cuda_module] +mod kernels { + use super::*; + + /// A static array through `shared_ptr`, a `from_raw` round trip, and a + /// negative `offset` back to this thread's slot. + #[kernel] + pub fn static_round_trip(out: *mut f32) { + let tid = thread::threadIdx_x() as usize; + // SAFETY: TILE_F32 is a `static mut SharedArray` of this kernel; each + // thread writes and reads its own slot, with a barrier between. + unsafe { + let tile = SharedArray::shared_ptr(&raw mut TILE_F32); + smem_fill(tile, tid, tid as f32 + 0.5); + thread::sync_threads(); + let again = SharedPtr::from_raw(tile.as_ptr()); + let last = again.add(THREADS - 1); + let mine = last.offset(-((THREADS - 1 - tid) as isize)); + *out.add(tid) = smem_load(mine, 0); + } + } + + /// A closure that captured a `SharedPtr`, called through a generic + /// non-inlined helper. + #[kernel] + pub fn closure_capture(out: *mut u32) { + let tid = thread::threadIdx_x() as usize; + // SAFETY: TILE_U32 is a `static mut SharedArray` of this kernel. + let tile = unsafe { SharedArray::shared_ptr(&raw mut TILE_U32) }; + // SAFETY: each thread writes its own slot. + let store = move |i: usize| unsafe { tile.add(i).as_ptr().write(i as u32 * 3) }; + apply(store, tid); + thread::sync_threads(); + // SAFETY: the slot was written before the barrier. + unsafe { *out.add(tid) = tile.add(tid).as_ptr().read() }; + } + + /// `addr_u32` deltas within the dynamic pool, and equality with + /// `cvta_generic_to_shared_u32` on the pool's raw pointer. + #[kernel] + pub fn addr_u32_deltas(out: *mut u32) { + let tid = thread::threadIdx_x() as usize; + let base = DynamicSharedArray::::shared_ptr(); + let at64 = DynamicSharedArray::::shared_ptr_at(64); + // SAFETY: `get_raw` points into the dynamic pool; thread 0 writes the + // four output slots. + unsafe { + let raw = cvta_generic_to_shared_u32(DynamicSharedArray::::get_raw()); + if tid == 0 { + *out.add(0) = at64.addr_u32() - base.addr_u32(); + *out.add(1) = base.add(3).addr_u32() - base.addr_u32(); + *out.add(2) = (base.addr_u32() == raw) as u32; + *out.add(3) = base.cast::().byte_add(7).addr_u32() - base.addr_u32(); + } + } + } + + /// `SharedPtr` behind references: `&Region` and `&mut SharedPtr`. + #[kernel] + pub fn behind_references(out: *mut u32) { + let tid = thread::threadIdx_x() as usize; + let pool = DynamicSharedArray::::shared_ptr(); + // SAFETY: the launch supplies 1024 dynamic shared bytes, enough for + // the 64 `u32` slots; every slot is written before the barrier. + unsafe { + pool.add(tid).as_ptr().write(100 + tid as u32); + thread::sync_threads(); + let region = Region { + base: pool, + offset: 0, + }; + let a = read_region(®ion, tid); + let mut cursor = pool; + bump(&mut cursor, tid); + let b = cursor.as_ptr().read(); + let shifted = Region { + base: pool, + offset: THREADS / 2, + }; + let c = read_region(&shifted, tid % (THREADS / 2)); + *out.add(tid) = a + b + c; + } + } + + /// A non-inlined helper that returns a `SharedPtr` picks between a static + /// array and the dynamic pool. + #[kernel] + pub fn returned_pointer(out: *mut f32) { + let tid = thread::threadIdx_x() as usize; + let dynp = DynamicSharedArray::::shared_ptr(); + // SAFETY: TILE_F32 is a `static mut SharedArray` of this kernel and + // the pool holds 64 `f32`; each thread writes its own slot in both. + unsafe { + let tile = SharedArray::shared_ptr(&raw mut TILE_F32); + smem_fill(tile, tid, 1000.0 + tid as f32); + smem_fill(dynp, tid, 2000.0 + tid as f32); + thread::sync_threads(); + let chosen = pick(tile, dynp, tid % 2 == 1); + *out.add(tid) = smem_load(chosen, tid); + } + } + + /// `[SharedPtr; 2]` passed by value to a non-inlined helper. + #[kernel] + pub fn array_by_value(out: *mut u32) { + let tid = thread::threadIdx_x() as usize; + let pool = DynamicSharedArray::::shared_ptr(); + // SAFETY: TILE_U32 is a `static mut SharedArray` of this kernel and + // the pool holds 64 `u32`; each thread writes its own slot in one. + unsafe { + let tile = SharedArray::shared_ptr(&raw mut TILE_U32); + store_via_array([tile, pool], tid % 2, tid, 7 + tid as u32); + thread::sync_threads(); + let v = if tid.is_multiple_of(2) { + tile.add(tid).as_ptr().read() + } else { + pool.add(tid).as_ptr().read() + }; + *out.add(tid) = v; + } + } + + /// A struct pointee, and a `cast` to its first field. + #[kernel] + pub fn struct_pointee(out: *mut f32) { + let tid = thread::threadIdx_x() as usize; + // SAFETY: PAIRS is a `static mut SharedArray` of this kernel; each + // thread writes and reads its own element. + unsafe { + let pairs = SharedArray::shared_ptr(&raw mut PAIRS); + smem_fill( + pairs, + tid, + Pair { + a: tid as f32, + b: 2.0 * tid as f32, + }, + ); + thread::sync_threads(); + let p = smem_load(pairs, tid); + let first: f32 = smem_load(pairs.add(tid).cast::(), 0); + *out.add(tid) = p.a + p.b + first; + } + } + + /// `byte_add` into a dynamic byte pool, viewed as `f64` and as `u16`. + #[kernel] + pub fn byte_views(out: *mut f64) { + let tid = thread::threadIdx_x() as usize; + let pool = DynamicSharedArray::::shared_ptr(); + // SAFETY: the pool holds 1024 bytes; the `f64` view uses the first + // 512 and the `u16` view the next 128, each thread its own slot. + unsafe { + let d = pool.byte_add(8 * tid).cast::(); + smem_fill(d, 0, tid as f64 * 0.25); + let h = pool.byte_add(8 * THREADS + 2 * tid).cast::(); + smem_fill(h, 0, tid as u16 + 7); + thread::sync_threads(); + let dv: f64 = smem_load(pool.cast::(), tid); + let hv: u16 = smem_load(pool.byte_add(8 * THREADS).cast::(), tid); + *out.add(tid) = dv + hv as f64; + } + } + + /// `Option>` returned from and passed into non-inlined + /// helpers. + #[kernel] + pub fn option_pointer(out: *mut u32) { + let tid = thread::threadIdx_x() as usize; + // SAFETY: TILE_U32 is a `static mut SharedArray` of this kernel; each + // thread writes and reads its own slot. + unsafe { + let tile = SharedArray::shared_ptr(&raw mut TILE_U32); + let took = store_opt(maybe(tile, true), tid, 5 + tid as u32); + let skipped = store_opt(maybe(tile, false), tid, 0); + thread::sync_threads(); + *out.add(tid) = tile.add(tid).as_ptr().read() + 10 * took + 100 * skipped; + } + } + + /// `SharedPtr` values stored in shared memory: a `SharedArray` of them. + #[kernel] + pub fn stored_in_shared(out: *mut u32) { + let tid = thread::threadIdx_x() as usize; + // SAFETY: TILE_U32 and SLOTS are `static mut SharedArray`s of this + // kernel; each thread writes its own slot in both and reads one slot + // after the barrier. + unsafe { + let tile = SharedArray::shared_ptr(&raw mut TILE_U32); + let slots = SharedArray::shared_ptr(&raw mut SLOTS); + tile.add(tid).as_ptr().write(tid as u32 * 11); + slots.add(tid).as_ptr().write(tile.add(THREADS - 1 - tid)); + thread::sync_threads(); + let p: SharedPtr = slots.add(tid).as_ptr().read(); + *out.add(tid) = p.as_ptr().read(); + } + } + + /// `transmute` to `usize` and back, in non-inlined helpers. The high word + /// records whether the bits of this thread's slot sit `4 * tid` bytes + /// past the bits of the array base. + #[kernel] + pub fn transmuted(out: *mut u64) { + let tid = thread::threadIdx_x() as usize; + // SAFETY: TILE_U32 is a `static mut SharedArray` of this kernel; each + // thread writes and reads its own slot. + unsafe { + let tile = SharedArray::shared_ptr(&raw mut TILE_U32); + let bits = to_bits(tile.add(tid)); + let back = from_bits(bits); + back.as_ptr().write(tid as u32 * 13); + thread::sync_threads(); + let v = tile.add(tid).as_ptr().read(); + let strided = bits == to_bits(tile) + 4 * tid; + *out.add(tid) = v as u64 | ((strided as u64) << 32); + } + } +} + +/// Counts the mismatches of one kernel and prints its line. +fn check(name: &str, bad: usize, mismatches: &mut usize, kernels: &mut usize) { + if bad == 0 { + println!(" ok {name}"); + } else { + println!(" BAD {name}: {bad} mismatches"); + } + *mismatches += bad; + *kernels += 1; +} + +fn main() -> Result<(), Box> { + let ctx = CudaContext::new(0)?; + let stream = ctx.default_stream(); + let module = kernels::load(&ctx)?; + let config = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (THREADS as u32, 1, 1), + shared_mem_bytes: 1024, + }; + let mut mismatches = 0usize; + let mut kernels = 0usize; + println!("=== SharedPtr API ({THREADS} threads) ==="); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + // SAFETY: one block of THREADS threads; every output buffer holds one + // slot per thread (four slots for addr_u32_deltas, written by thread 0). + unsafe { module.static_round_trip(&stream, config, out.cu_deviceptr() as *mut f32)? }; + let got = out.to_host_vec(&stream)?; + check( + "static_round_trip", + (0..THREADS).filter(|&i| got[i] != i as f32 + 0.5).count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.closure_capture(&stream, config, out.cu_deviceptr() as *mut u32)? }; + let got = out.to_host_vec(&stream)?; + check( + "closure_capture", + (0..THREADS).filter(|&i| got[i] != i as u32 * 3).count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, 4)?; + unsafe { module.addr_u32_deltas(&stream, config, out.cu_deviceptr() as *mut u32)? }; + let got = out.to_host_vec(&stream)?; + let want = [64u32, 12, 1, 7]; + check( + "addr_u32_deltas", + (0..4).filter(|&i| got[i] != want[i]).count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.behind_references(&stream, config, out.cu_deviceptr() as *mut u32)? }; + let got = out.to_host_vec(&stream)?; + check( + "behind_references", + (0..THREADS) + .filter(|&i| { + let t = i as u32; + let half = (THREADS / 2) as u32; + got[i] != (100 + t) + (100 + t) + (100 + half + t % half) + }) + .count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.returned_pointer(&stream, config, out.cu_deviceptr() as *mut f32)? }; + let got = out.to_host_vec(&stream)?; + check( + "returned_pointer", + (0..THREADS) + .filter(|&i| { + let base = if i % 2 == 1 { 2000.0 } else { 1000.0 }; + got[i] != base + i as f32 + }) + .count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.array_by_value(&stream, config, out.cu_deviceptr() as *mut u32)? }; + let got = out.to_host_vec(&stream)?; + check( + "array_by_value", + (0..THREADS).filter(|&i| got[i] != 7 + i as u32).count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.struct_pointee(&stream, config, out.cu_deviceptr() as *mut f32)? }; + let got = out.to_host_vec(&stream)?; + check( + "struct_pointee", + (0..THREADS).filter(|&i| got[i] != 4.0 * i as f32).count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.byte_views(&stream, config, out.cu_deviceptr() as *mut f64)? }; + let got = out.to_host_vec(&stream)?; + check( + "byte_views", + (0..THREADS) + .filter(|&i| got[i] != i as f64 * 0.25 + (i as f64 + 7.0)) + .count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.option_pointer(&stream, config, out.cu_deviceptr() as *mut u32)? }; + let got = out.to_host_vec(&stream)?; + check( + "option_pointer", + (0..THREADS).filter(|&i| got[i] != 15 + i as u32).count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.stored_in_shared(&stream, config, out.cu_deviceptr() as *mut u32)? }; + let got = out.to_host_vec(&stream)?; + check( + "stored_in_shared", + (0..THREADS) + .filter(|&i| got[i] != (THREADS - 1 - i) as u32 * 11) + .count(), + &mut mismatches, + &mut kernels, + ); + + let out = DeviceBuffer::::zeroed(&stream, THREADS)?; + unsafe { module.transmuted(&stream, config, out.cu_deviceptr() as *mut u64)? }; + let got = out.to_host_vec(&stream)?; + check( + "transmuted", + (0..THREADS) + .filter(|&i| got[i] != (i as u64 * 13) | (1u64 << 32)) + .count(), + &mut mismatches, + &mut kernels, + ); + + if mismatches == 0 { + println!("SUCCESS: shared_ptr_api, {kernels} kernels, 0 mismatches"); + Ok(()) + } else { + println!("FAIL: shared_ptr_api, {mismatches} mismatches across {kernels} kernels"); + std::process::exit(1) + } +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.lock index 6b72183413..098e6fa943 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.toml index 45303244c7..2c87a73b3d 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_ptr_transmute/Cargo.toml @@ -23,4 +23,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 6aa62f47e1..bc9bbac296 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 @@ -30,6 +30,12 @@ //! the integer metadata while the data-pointer half follows the compiler's //! address-space rules. The cast must remain an SSA reconstruction rather //! than an aggregate memory round-trip. +//! - union view: a union holding the shared pointer next to a `u64` stores +//! the pointer as a generic pointer, so its integer view must equal the +//! transmuted `usize` above, the bits must recover the same pointer, and +//! a store through that pointer must be visible through the allocation. +//! The union is built and read by value; in-place writes to the pointer +//! field are refused by the compiler. //! //! Each scalar transmute and the fat-pointer cast live in `#[inline(never)]` //! device functions so LLVM cannot see both halves of a round trip in one body @@ -49,6 +55,16 @@ mod kernels { const THREADS: usize = 32; + /// One byte image, two views: the shared pointer as the program sees it, + /// and the same 8 bytes as an integer. The compiler stores the pointer as + /// a generic pointer, so the image is the same on every target even + /// though a shared-memory pointer is 4 bytes wide under modern NVVM. + #[repr(C)] + union SharedPtrBits { + ptr: *mut SharedArray, + bits: u64, + } + // Clippy would rewrite these transmutes as `as` casts, but the `as` // form lowers through a different compiler path (PtrToAddress / // pointer casts); `Transmute` lowering is the path under test here. @@ -158,6 +174,22 @@ mod kernels { *out.get_unchecked_mut(5) = core::ptr::eq(recast_data.cast::(), base) as u32; recast_data.add(1).write(0x1234_i32); *out.get_unchecked_mut(6) = SCRATCH[1]; + + // Union view, by value in both directions. The integer view + // is the Rust-visible address (the transmute result above), + // the bits recover the pointer, and a store through the + // recovered pointer lands in the shared allocation. The + // store is guarded by the pointer check so a wrong pointer + // is reported as a failure instead of faulting. + let union_bits = SharedPtrBits { ptr: raw }.bits; + let from_bits = SharedPtrBits { bits: bits as u64 }.ptr; + let from_bits_eq = core::ptr::eq(from_bits, raw); + *out.get_unchecked_mut(7) = (union_bits == bits as u64) as u32; + *out.get_unchecked_mut(8) = from_bits_eq as u32; + if from_bits_eq { + (&mut (*from_bits))[2] = 0xBEEF; + } + *out.get_unchecked_mut(9) = SCRATCH[2]; } } } @@ -175,11 +207,11 @@ fn main() -> Result<(), Box> { block_dim: (32, 1, 1), shared_mem_bytes: 0, }; - let mut out = DeviceBuffer::::zeroed(&stream, 7)?; + let mut out = DeviceBuffer::::zeroed(&stream, 10)?; // SAFETY: one 32-thread block writes 32 distinct shared elements, then // thread 0 reads or mutates the shared allocation after a block-wide - // barrier. Only thread 0 writes the seven output elements. + // barrier. Only thread 0 writes the ten output elements. unsafe { module.shared_ptr_transmute(&stream, cfg, &mut out) }?; let result = out.to_host_vec(&stream)?; @@ -224,9 +256,28 @@ fn main() -> Result<(), Box> { ); std::process::exit(1); } + if result[7] != 1 { + eprintln!( + "FAIL shared_ptr_transmute: shared-pointer union integer view differs from the pointer's address bits" + ); + std::process::exit(1); + } + if result[8] != 1 { + eprintln!( + "FAIL shared_ptr_transmute: shared-pointer union bits did not recover the pointer" + ); + std::process::exit(1); + } + if result[9] != 0xBEEF { + eprintln!( + "FAIL shared_ptr_transmute: store through the union-recovered pointer read back {:#X}, expected 0xBEEF", + result[9] + ); + std::process::exit(1); + } println!( - "PASS shared_ptr_transmute: scalar and slice-fat-pointer casts preserve shared-pointer semantics, sum={}, scalar_readback={:#X}, slice_readback={:#X}", - result[1], result[3], result[6] + "PASS shared_ptr_transmute: scalar, slice-fat-pointer and union views preserve shared-pointer semantics, sum={}, scalar_readback={:#X}, slice_readback={:#X}, union_readback={:#X}", + result[1], result[3], result[6], result[9] ); Ok(()) } 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 0a11acde77..fd1d381de4 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/shared_slice_unsize/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/Cargo.toml index 358296538a..e16dad84e6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_slice_unsize/Cargo.toml @@ -22,4 +22,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 b3f2de5290..5ac0486dc4 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/shared_stage_select/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/Cargo.toml index 2d74c93a6d..e9aabd0542 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shared_stage_select/Cargo.toml @@ -16,4 +16,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 e9f87595fa..2bee1e9097 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.toml index f093c77424..661a46e963 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime (with HMM support) -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/README.md index 4533fdf4f6..269a200951 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sharedmem/README.md @@ -8,6 +8,10 @@ Demonstrates `SharedArray` for block-level cooperation. Shared memory is ~ 1. **shared_test**: Single SharedArray, each thread reads neighbor's value 2. **shared_dual**: Two SharedArrays, threads combine neighbor values from both +3. **shared_u64_align1**: A `SharedArray` asks for 1-byte alignment; its storage is still declared at 8 bytes +4. **shared_v4_natural**: A `SharedArray` with a `#[repr(C, align(16))]` element and no `ALIGN` is declared at 16 bytes + +The host checks every kernel's values, then reads `sharedmem.ptx` to confirm the two `.align` declarations. ## Key Concepts Demonstrated @@ -53,6 +57,29 @@ static mut TILE_B: SharedArray = SharedArray::UNINIT; Each `static mut` creates a separate shared memory allocation. +### Storage Alignment + +Every shared array is declared at the larger of the requested `ALIGN` and the +element's ABI alignment. `ALIGN` must be 0 (natural) or a power of two; any +other value is a compile error that names the use site. + +```rust +#[repr(C, align(16))] +#[derive(Copy, Clone)] +pub struct V4(pub [f32; 4]); + +static mut DATA: SharedArray = SharedArray::UNINIT; // declared .align 8 +static mut VS: SharedArray = SharedArray::UNINIT; // declared .align 16 +``` + +Both alignment kernels touch the under-requested array before a smaller pad +array. Shared globals are inserted at the front of the module, so the first +array a kernel references lands last in the PTX, right behind the pad. +Declared at the requested 1 byte, or at the old blanket 8 for `V4`, ptxas +packed the array at an offset its own `st.shared.b64` and `st.shared.v4.b32` +could not use, and the launch failed with CUDA error 716, `misaligned +address`. With the ABI raise the arrays sit on 8- and 16-byte boundaries. + ## Build and Run ```bash @@ -75,6 +102,17 @@ Input b[0..5] = [100.0, 101.0, 102.0, 103.0, 104.0] Output out[0..5] = [102.0, 104.0, 106.0, 108.0, 110.0] ✓ Dual SharedArray: correct neighbor sum from both tiles +=== Test 3: SharedArray behind a byte pad === +Output out[0..4] = [1122334455667788, 112233445566778a, 112233445566778a, 112233445566778c] +✓ u64 storage requested at 1 byte reads back intact + +=== Test 4: SharedArray behind a u64 pad === +Output out[0..4] = [3.5, 6.5, 9.5, 12.5] +✓ align(16) element at natural alignment reads back intact + +=== Test 5: PTX storage declarations === +✓ PTX declares the u64 array at .align 8 and the V4 array at .align 16 + ✓ SUCCESS: All shared memory tests passed! ``` @@ -140,6 +178,10 @@ for stride in [128, 64, 32, 16, 8, 4, 2, 1] { // Shared memory declaration .shared .align 4 .b8 TILE[1024]; // 256 * sizeof(f32) +// Storage declared at the element's ABI alignment, whatever ALIGN asked for +.shared .align 8 .b8 DATA[16]; // SharedArray +.shared .align 16 .b8 VS[64]; // SharedArray + // Write st.shared.f32 [%rd_tile + %r_offset], %f_val; 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 6ae34ad7c6..7852b13c0f 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 @@ -11,6 +11,15 @@ //! Demonstrates SharedArray for block-level cooperation: //! 1. shared_test - single SharedArray: loads data, syncs, reads neighbor //! 2. shared_dual - two SharedArrays: loads from a and b, syncs, adds neighbors +//! 3. shared_u64_align1 - a `u64` array requested at 1-byte alignment is +//! still declared at 8 bytes +//! 4. shared_v4_natural - a `#[repr(C, align(16))]` element at natural +//! alignment is declared at 16 bytes +//! +//! The last two are regressions for the storage alignment contract: every +//! shared array is declared at the larger of the requested `ALIGN` and the +//! element's ABI alignment. The host checks the kernel values and the +//! `.align` declarations in the generated PTX. //! //! Build and run with: //! cargo oxide run sharedmem @@ -20,6 +29,19 @@ use cuda_core::{CudaContext, DeviceBuffer}; use cuda_device::{DisjointSlice, SharedArray, kernel, thread}; use cuda_host::cuda_module; +/// A 16-byte aligned aggregate. The LLVM struct type cannot carry the +/// `align(16)` raise; it lives on the Rust type and must reach the storage +/// declaration. +#[repr(C, align(16))] +#[derive(Copy, Clone)] +pub struct V4(pub [f32; 4]); + +/// Marker value written by the `u64` alignment kernel. +const U64_BASE: u64 = 0x1122_3344_5566_7788; + +/// Threads in the block of the two alignment kernels. +const ALIGN_THREADS: usize = 32; + // ============================================================================= // KERNELS // ============================================================================= @@ -77,12 +99,103 @@ mod kernels { } } } + + /// `DATA` asks for 1-byte alignment but holds `u64` values. + /// + /// It is touched before `PAD` on purpose. Shared globals are inserted at + /// the front of the module, so the first array a kernel references lands + /// last in the PTX, right behind the 33-byte pad. Declared at 1 byte, + /// ptxas packs it at an odd offset and `st.shared.b64` faults with + /// `misaligned address`. Declared at the `u64` ABI alignment, it is + /// placed on an 8-byte boundary. + #[kernel] + pub fn shared_u64_align1(mut out: DisjointSlice) { + static mut DATA: SharedArray = SharedArray::UNINIT; + static mut PAD: SharedArray = SharedArray::UNINIT; + + let tid = thread::threadIdx_x() as usize; + unsafe { + if tid < 2 { + DATA[tid] = U64_BASE + tid as u64; + } + if tid < 33 { + PAD[tid] = tid as u8; + } + } + + thread::sync_threads(); + + unsafe { + if let Some(out_elem) = out.get_mut(thread::index_1d()) { + *out_elem = DATA[tid % 2] + PAD[tid % 33] as u64; + } + } + } + + /// `VS` holds a `#[repr(C, align(16))]` element at natural alignment. + /// + /// Same ordering as `shared_u64_align1`: `VS` is touched before the + /// 8-byte `PAD`, so it lands behind the pad in the PTX. Its loads and + /// stores are 16-byte vector accesses, so the storage must be declared + /// at 16 bytes even though no `ALIGN` was requested. + #[kernel] + pub fn shared_v4_natural(mut out: DisjointSlice) { + static mut VS: SharedArray = SharedArray::UNINIT; + static mut PAD: SharedArray = SharedArray::UNINIT; + + let tid = thread::threadIdx_x() as usize; + let t = tid as f32; + unsafe { + if tid < 4 { + VS[tid] = V4([t, t + 0.5, -t, 2.0 * t]); + } + if tid == 0 { + PAD[0] = 3; + } + } + + thread::sync_threads(); + + unsafe { + let v = VS[tid % 4]; + if let Some(out_elem) = out.get_mut(thread::index_1d()) { + *out_elem = v.0[0] + v.0[1] + v.0[2] + v.0[3] + PAD[0] as f32; + } + } + } } // ============================================================================= // HOST CODE // ============================================================================= +/// Check the two declarations the alignment kernels rely on in the PTX that +/// `cargo oxide` writes next to this manifest. The byte sizes identify the +/// arrays: `DATA` is 16 bytes and `VS` is 64 bytes. +fn verify_shared_alignment_ptx() -> Result<(), Box> { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("sharedmem.ptx"); + let ptx = std::fs::read_to_string(&path)?; + for (align, bytes, array) in [ + (8, 16, "SharedArray"), + (16, 64, "SharedArray"), + ] { + let directive = format!(".shared .align {align} .b8 "); + let extent = format!("[{bytes}];"); + let count = ptx + .lines() + .filter(|line| line.contains(&directive) && line.trim_end().ends_with(&extent)) + .count(); + if count != 1 { + return Err(format!( + "{}: expected one `{directive}...{extent}` declaration for {array}, found {count}", + path.display() + ) + .into()); + } + } + Ok(()) +} + fn main() { println!("=== Unified Shared Memory Example ===\n"); @@ -167,5 +280,67 @@ fn main() { println!("✓ Dual SharedArray: correct neighbor sum from both tiles\n"); } + // One block of ALIGN_THREADS threads for the two alignment kernels. + let align_cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (ALIGN_THREADS as u32, 1, 1), + shared_mem_bytes: 0, + }; + + // ===== Test 3: u64 array requested at 1-byte alignment ===== + println!("=== Test 3: SharedArray behind a byte pad ==="); + { + let mut out_dev = DeviceBuffer::::zeroed(&stream, ALIGN_THREADS).unwrap(); + + // SAFETY: one block of ALIGN_THREADS threads; the buffer holds one element per thread. + unsafe { module.shared_u64_align1(stream.as_ref(), align_cfg, &mut out_dev) } + .expect("Kernel launch failed"); + + let out_result = out_dev.to_host_vec(&stream).unwrap(); + println!("Output out[0..4] = {:x?}", &out_result[0..4]); + + // Verify: out[i] = DATA[i % 2] + PAD[i % 33] + for (i, &got) in out_result.iter().enumerate() { + let expected = U64_BASE + (i % 2) as u64 + (i % 33) as u64; + if got != expected { + eprintln!("Mismatch at {i}: expected {expected:#x}, got {got:#x}"); + std::process::exit(1); + } + } + println!("✓ u64 storage requested at 1 byte reads back intact\n"); + } + + // ===== Test 4: repr(align(16)) element at natural alignment ===== + println!("=== Test 4: SharedArray behind a u64 pad ==="); + { + let mut out_dev = DeviceBuffer::::zeroed(&stream, ALIGN_THREADS).unwrap(); + + // SAFETY: one block of ALIGN_THREADS threads; the buffer holds one element per thread. + unsafe { module.shared_v4_natural(stream.as_ref(), align_cfg, &mut out_dev) } + .expect("Kernel launch failed"); + + let out_result = out_dev.to_host_vec(&stream).unwrap(); + println!("Output out[0..4] = {:?}", &out_result[0..4]); + + // Verify: VS[s] sums to 3s + 0.5 for s = i % 4, plus PAD[0] = 3. + for (i, &got) in out_result.iter().enumerate() { + let s = (i % 4) as f32; + let expected = 3.0 * s + 0.5 + 3.0; + if got != expected { + eprintln!("Mismatch at {i}: expected {expected}, got {got}"); + std::process::exit(1); + } + } + println!("✓ align(16) element at natural alignment reads back intact\n"); + } + + // ===== Test 5: PTX declarations ===== + println!("=== Test 5: PTX storage declarations ==="); + if let Err(err) = verify_shared_alignment_ptx() { + eprintln!("PTX declaration check failed: {err}"); + std::process::exit(1); + } + println!("✓ PTX declares the u64 array at .align 8 and the V4 array at .align 16\n"); + println!("✓ SUCCESS: All shared memory tests passed!"); } 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 7e138e6546..9c58993282 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/shuffle_64/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/Cargo.toml index a6dc731fa9..e73fe371c3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/shuffle_64/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 205c918587..dcfee5f6a1 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/Cargo.toml index c8c2483561..abeaec12f6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_from_raw_parts_cast/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.lock index 9edf3c367b..c94c1a3938 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.toml index 443f9a35cc..eaf98d3b20 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_get_mut/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 5f89a85bdc..a4045435e3 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/Cargo.toml index 8a79c5c642..ef9ea202ef 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_index_bounds_check/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 a0b6344167..6483648ff7 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/slice_iterators/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/Cargo.toml index edecdd9f2a..fd68089b36 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_iterators/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.lock index 1a2c8f2fad..44e93548de 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.toml index 495beb784e..277283ffc3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/slice_reslice/Cargo.toml @@ -21,4 +21,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 4f6f48790e..20196ba508 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/small_type_ffi_test/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/Cargo.toml index 09bdec2a73..eab7e6e578 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/small_type_ffi_test/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 9b2a168a17..8b7f279ddd 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/Cargo.toml index e34849d2ca..6788749c8b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_fp8/Cargo.toml @@ -13,4 +13,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 7e0f14c7f6..45c9b4ba13 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 @@ -3,8 +3,8 @@ //! Checks all four plain SM89 sparse FP8 MMA forms against an exact host GEMM. //! -//! Each form covers all twelve standard-metadata codes and zero/nonzero C. -//! Values vary across K bands so swapped fragment registers change the result. +//! Each form covers all six valid FP8 metadata codes and zero/nonzero C. +//! Values vary across K bands so swapped B fragment registers change the result. //! See the PTX ISA's "Matrix Fragments for sparse mma.m16n8k64" for the layout. //! `--probe` additionally compares candidate layouts on the current GPU. @@ -15,7 +15,7 @@ use cuda_device::{DisjointSlice, cuda_module, kernel, thread, wmma}; const M: usize = 16; const N: usize = 8; const K: usize = 64; -const METADATA_CODES: usize = 12; +const METADATA_CODES: usize = 6; const VARIANTS: usize = 2 * 4 * METADATA_CODES; /// Compressed K: two kept elements per 4-wide group. const KC: usize = K / 2; @@ -59,9 +59,19 @@ fn f8_bits(v: u32, e4m3: bool) -> u32 { } } -/// Standard metadata accepts two distinct 2-bit indices, in either order. +/// Plain FP8 `mma.sp` requires two positions in increasing order, even without +/// `::ordered_metadata`: `0x4` encodes (0, 1), but reversed `0x1` is undefined. +/// See PTX 8.4's FP8-specific rule under "Sparse matrix storage": +/// . fn nibble_of(i: u32) -> u32 { - probe_nibble(i) + match i { + 0 => 0x4, + 1 => 0x8, + 2 => 0x9, + 3 => 0xc, + 4 => 0xd, + _ => 0xe, + } } #[cuda_module] @@ -132,7 +142,7 @@ mod kernels { } /// Variant `v` occupies `out[v*128 + lane*4 .. +4]`, with - /// `v = nonzero_accumulator*48 + form*12 + metadata_code`. + /// `v = (nonzero_accumulator*4 + form)*METADATA_CODES + metadata_code`. #[kernel] pub fn oracle(mut out: DisjointSlice) { let l = thread::threadIdx_x(); @@ -211,14 +221,9 @@ mod kernels { let mut mi = 0u32; while mi < META_OPTS { - let nib = probe_nibble(mi % 12); - // `sw` exchanges the two 2-bit fields of the nibble. - let word = if mi / 12 == 0 { - nib - } else { - ((nib & 3) << 2) | (nib >> 2) - }; - let meta = word * 0x1111_1111; + // Vary fragment layouts, but never reverse the metadata's + // two indices: those FP8 inputs have undefined behavior. + let meta = nibble_of(mi) * 0x1111_1111; let d = unsafe { wmma::mma_sp_m16n8k64_f32_e4m3_e4m3_f32(zero, a, b, meta, 0) }; let p = combo * 128 + (l as usize) * 4; let mut j = 0usize; @@ -292,7 +297,7 @@ fn run_oracle(ctx: &std::sync::Arc) { } } println!( - " {name:12} {:>7} all 12 standard metadata codes: {} mismatches", + " {name:12} {:>7} all {METADATA_CODES} FP8 metadata codes: {} mismatches", if acc == 1 { "C!=0," } else { "C=0," }, form_bad ); @@ -301,7 +306,7 @@ fn run_oracle(ctx: &std::sync::Arc) { } assert_eq!(bad, 0, "sparse FP8 MMA accumulator mismatches"); println!( - "SUCCESS: all 4 sparse FP8 m16n8k64 forms x 12 metadata codes x C=0/C!=0; \ + "SUCCESS: all 4 sparse FP8 m16n8k64 forms x {METADATA_CODES} metadata codes x C=0/C!=0; \ all 32 lanes and 4 logical accumulators/lane match host GEMM exactly" ); } @@ -316,12 +321,12 @@ fn run_oracle(ctx: &std::sync::Arc) { // A order `byte_ord(ao, i)` -- byte position i -> column offset // B band `b_band(bb, n)` -- which dense K rows a register holds // B order `byte_ord(bo, i)` -- byte position i -> K offset -// code the nibble and which of its two 2-bit fields names the even column +// code one of the six valid FP8 metadata nibbles // // Every candidate is filled into the registers and executed once; the host -// compares the result with an exact GEMM whose answer does not depend on the -// candidate at all. A wrong axis changes the sum, so only the true wiring -// reproduces it. +// compares the result with a GEMM using that candidate's layout and metadata. +// Exact matches are diagnostic candidates: repeated input patterns can leave +// more than one layout indistinguishable. // --------------------------------------------------------------------------- const A_ROW_OPTS: u32 = 4; @@ -329,11 +334,9 @@ const A_BAND_OPTS: u32 = 6; const A_ORD_OPTS: u32 = 2; const B_BAND_OPTS: u32 = 6; const B_ORD_OPTS: u32 = 2; -const META_SWAP_OPTS: u32 = 2; - const A_LAYOUTS: u32 = A_ROW_OPTS * A_BAND_OPTS * A_ORD_OPTS; const B_LAYOUTS: u32 = B_BAND_OPTS * B_ORD_OPTS; -const META_OPTS: u32 = 12 * META_SWAP_OPTS; +const META_OPTS: u32 = METADATA_CODES as u32; const COMBOS: u32 = A_LAYOUTS * B_LAYOUTS * META_OPTS; /// Separate K-band patterns for the optional layout sweep. @@ -392,35 +395,6 @@ fn byte_ord(o: u32, i: u32) -> u32 { if o == 0 { i } else { 3 - i } } -/// The twelve legal nibbles: two 2-bit fields that differ. -fn probe_nibble(idx: u32) -> u32 { - if idx == 0 { - 0x1 - } else if idx == 1 { - 0x2 - } else if idx == 2 { - 0x3 - } else if idx == 3 { - 0x4 - } else if idx == 4 { - 0x6 - } else if idx == 5 { - 0x7 - } else if idx == 6 { - 0x8 - } else if idx == 7 { - 0x9 - } else if idx == 8 { - 0xb - } else if idx == 9 { - 0xc - } else if idx == 10 { - 0xd - } else { - 0xe - } -} - #[derive(Clone, Copy)] struct Candidate { rp: u32, @@ -429,7 +403,6 @@ struct Candidate { bb: u32, bo: u32, nib: u32, - sw: u32, } impl Candidate { @@ -471,11 +444,7 @@ impl Candidate { fn kmap(&self, c: u32) -> u32 { let low = self.nib & 3; let high = self.nib >> 2; - let field = if c.is_multiple_of(2) == (self.sw == 0) { - low - } else { - high - }; + let field = if c.is_multiple_of(2) { low } else { high }; 4 * (c / 2) + field } @@ -546,8 +515,7 @@ fn run_probe(ctx: &std::sync::Arc) { ao: ai % A_ORD_OPTS, bb: bi / B_ORD_OPTS, bo: bi % B_ORD_OPTS, - nib: probe_nibble(mi % 12), - sw: mi / 12, + nib: nibble_of(mi), }; let idx = total; total += 1; @@ -567,8 +535,8 @@ fn run_probe(ctx: &std::sync::Arc) { for c in &exact { let (ap, bp) = c.coverage(); println!( - " exact rp={} ab={} ao={} bb={} bo={} nib={:#x} sw={} (A {ap}/512, B {bp}/512 placed)", - c.rp, c.ab, c.ao, c.bb, c.bo, c.nib, c.sw + " exact rp={} ab={} ao={} bb={} bo={} nib={:#x} (A {ap}/512, B {bp}/512 placed)", + c.rp, c.ab, c.ao, c.bb, c.bo, c.nib ); } } @@ -607,15 +575,18 @@ mod tests { } #[test] - fn metadata_covers_every_distinct_pair() { + fn metadata_covers_every_valid_fp8_pair() { let mut seen = [false; 16]; for index in 0..METADATA_CODES as u32 { let code = nibble_of(index) as usize; - assert_ne!(code & 3, code >> 2); assert!(!seen[code]); seen[code] = true; } - assert_eq!(seen.iter().filter(|&&present| present).count(), 12); + // FP8 requires the first stored value's position to precede the second. + // This checks both that reversed pairs are absent and no valid pair is lost. + for (code, present) in seen.into_iter().enumerate() { + assert_eq!(present, (code & 3) < (code >> 2), "metadata {code:#x}"); + } } #[test] 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 337858588b..218f65f807 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/Cargo.toml index 448364d3a1..b2271ca823 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/sparse_mma_ordered_float/Cargo.toml @@ -13,4 +13,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 60eb51980e..4fa006d2b4 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/split_at_mut/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/Cargo.toml index 3c4c198811..8e09a83268 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/split_at_mut/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/Cargo.lock index bd369b09ad..4870b6e0d2 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -715,3 +745,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.lock index 651695d697..260515c3a0 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.toml index c0036b8fd8..c28a14dd7b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_addend/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.lock index 3b9c78d986..d24fdb64fe 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.toml index 952d113a58..2e5350714b 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/static_slice_unsize/Cargo.toml @@ -12,6 +12,6 @@ license = "Apache-2.0" # cargo oxide run static_slice_unsize [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.lock index 8433420962..1b1df6b95f 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/step_by/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.toml index 764a8c0f47..944098be1c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/step_by/Cargo.toml @@ -22,4 +22,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 00a570174e..481d2d5496 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/struct_constant_provenance/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/Cargo.toml index 1c7b19f7da..e54f87c149 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_constant_provenance/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.lock index 3fdd686f8a..6e0c3e8691 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.toml index eacf1caf30..7fe0ac7e76 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/struct_field_layout/Cargo.toml @@ -20,4 +20,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 b4cc31859f..e32ae7b8bf 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/subslice_projection/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/Cargo.toml index abdea46862..69db8c6b33 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/README.md index 476d1b693c..9c090d6afa 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/subslice_projection/README.md @@ -7,11 +7,13 @@ Rustc emits `Subslice` for the `middle @ ..` portion of array and slice patterns - arrays use `Subslice { from, to, from_end: false }`; the result is the sized array place `[T; to - from]` starting at element `from`; - slices use `Subslice { from, to, from_end: true }`; the result keeps a data pointer advanced by `from` elements and metadata `old_len - from - to`. -Before this fix, the place walkers under `cuda-oxide/crates/mir-importer/src/translator/rvalue/` had no `Subslice` lowering. The value walker reported `Projection element ... not yet implemented in iterative mode`, while the address walker returned `Ok(None)`. Mutable borrows then failed loudly because falling back to a reference to a copy would lose write-through semantics. +`Subslice` lowering already supports sized arrays and slice fat pointers. This example also covers `ConstantIndex { from_end: true }` immediately following a slice `Subslice`. The index must use the rebuilt subslice metadata: `subslice_len = old_len - from - to`, then `index = subslice_len - offset`. Coverage includes both value reads and address-based mutable writes, plus a field projection after the from-end index. + +A slice `Subslice` that does not directly follow a fat-slice deref is still unsupported and fails loudly on purpose: a nested `[_, middle @ .., _]` then `[_, inner @ .., _]`, or a `Subslice` on a slice-tail field, stops with `slice Subslice reached iterative lowering without preserved fat-pointer metadata`. ## Coverage -The example contains five cases: +The example contains eight cases: | Case | MIR property checked | |---|---| @@ -20,14 +22,20 @@ The example contains five cases: | `array mutable ref` | mutable reference writes through to original array storage | | `slice metadata` | slice data pointer advances and length becomes `old_len - from - to` | | `slice mutable ref` | mutable slice subslice writes through to original storage | +| `slice subslice from-end field` | rebuilt slice metadata drives a from-end index followed by struct field projections | +| `slice subslice from-end mutable` | two distinct from-end offsets write through the rebuilt subslice while unrelated elements remain unchanged | +| `slice subslice from-end runtime` | a slice length read from the row data (3 to 8) drives two from-end writes and one from-end read on a local buffer; the host checks the packed result with the same helpers | -The helper functions are `#[inline(never)]` so optimized MIR retains the relevant projection in a separate body. +The helper functions are `#[inline(never)]` so optimized MIR retains the relevant projection in a separate body. LLVM still inlines them into the kernels, and the four-element cases then fold the from-end index to a constant element offset, so only the runtime-length case keeps the from-end arithmetic in the PTX: its accesses are `[%rd+-8]` for `[.., last]` and `[%rd+-12]` for `[.., second_last, _]`, off `base + (n << 2)` and behind the pattern-length guards. LLVM folds the `len - offset` subtraction into the displacement, so the PTX carries no `sub` instruction for it. Typical MIR shapes are expected to include: ```text Subslice { from: 1, to: 3, from_end: false } # [u32; 4] -> [u32; 2] Subslice { from: 1, to: 1, from_end: true } # [u32] -> [u32] +Subslice -> ConstantIndex { offset: 1, from_end: true } +Subslice -> ConstantIndex { offset: 2, from_end: true } +Subslice -> ConstantIndex { offset: 1, from_end: true } -> Field ``` ## Run 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 a156a8f26e..f05dd15c0c 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 @@ -13,7 +13,9 @@ * metadata must become `old_len - from - to`. * * The helpers are `#[inline(never)]` so optimized MIR keeps the projection - * in the callee instead of dissolving it into the kernel. + * in the callee instead of dissolving it into the kernel. One case slices a + * local buffer at a lane-dependent length, so the PTX keeps the from-end + * arithmetic that the constant-length cases fold away. * * Build/run: * cargo oxide run subslice_projection @@ -26,6 +28,23 @@ use cuda_device::{DisjointSlice, kernel, thread}; use cuda_host::cuda_module; use std::sync::Arc; +#[derive(Copy, Clone)] +struct Pair { + x: u32, + y: u32, +} + +#[inline(never)] +fn middle_last_field(values: &[Pair]) -> u32 { + let [_, middle @ .., _] = values else { + return 0; + }; + let [.., last] = middle else { + return 0; + }; + last.x + last.y +} + #[inline(never)] fn array_middle_value(values: [u32; 4]) -> [u32; 2] { let [_, middle @ .., _] = values; @@ -53,10 +72,120 @@ fn bump_slice_middle(values: &mut [u32]) { } } +// The two-step pattern is deliberate: `[_, middle @ .., _]` followed by +// `[.., last]` makes rustc emit a from-end `ConstantIndex` after a slice +// `Subslice`, `(*_1)[1:-1][-1 of 1]`. The single pattern `[_, .., last, _]` +// would index the original slice directly and skip the path under test. +#[inline(never)] +fn bump_slice_middle_last(values: &mut [u32]) { + if let [_, middle @ .., _] = values + && let [.., last] = middle + { + *last += 11; + } +} + +#[inline(never)] +fn bump_slice_middle_second_last(values: &mut [u32]) { + if let [_, middle @ .., _] = values + && let [.., second_last, _] = middle + { + *second_last += 13; + } +} + +#[inline(never)] +fn middle_last(values: &[u32]) -> u32 { + let [_, middle @ .., _] = values else { + return 0; + }; + let [.., last] = middle else { + return 0; + }; + *last +} + +/// Elements in the local buffer of the runtime-length case. +const BUFFER_LEN: usize = 8; + +/// Slice length for the runtime-length case. It comes from the row data, which +/// the compiler cannot see, so the from-end index stays a runtime computation. +/// `make_input` puts `10 + i` in column 0, so lane `i` works on `3 + i` +/// elements: from a one-element middle, where only `[.., last]` matches, up to +/// the whole buffer. +fn runtime_len(row: [u32; 4]) -> usize { + (row[0] - 7) as usize +} + +/// Seeds the local buffer from `row`, runs the from-end helpers on its first +/// `n` elements (`n <= BUFFER_LEN`) and packs the outcome into one word: every +/// element's delta weighted by its 1-based position, times 1000, plus the +/// value `middle_last` reads back. A write to the wrong element, a collateral +/// write and a read from the wrong element each change the word. The kernel +/// and the host oracle both call this function, so the device result is +/// checked against the same helpers compiled for the host. +fn from_end_at_runtime_len(row: [u32; 4], n: usize) -> u32 { + let mut values: [u32; BUFFER_LEN] = [ + row[0], + row[1], + row[2], + row[3], + row[0] + 50, + row[1] + 50, + row[2] + 50, + row[3] + 50, + ]; + let before = values; + + let window = &mut values[..n]; + bump_slice_middle_last(window); + bump_slice_middle_second_last(window); + let read = middle_last(&values[..n]); + + let mut weighted = 0; + for k in 0..values.len() { + weighted += (values[k] - before[k]) * (k as u32 + 1); + } + weighted * 1000 + read +} + #[cuda_module] mod kernels { use super::*; + #[kernel] + pub fn test_slice_subslice_from_end_field(input: &[[u32; 4]], mut out: DisjointSlice) { + let idx = thread::index_1d(); + let i = idx.get(); + if i >= input.len() { + return; + } + + let row = input[i]; + let values = [ + Pair { + x: row[0], + y: row[0] + 100, + }, + Pair { + x: row[1], + y: row[1] + 100, + }, + Pair { + x: row[2], + y: row[2] + 100, + }, + Pair { + x: row[3], + y: row[3] + 100, + }, + ]; + + if let Some(slot) = out.get_mut(idx) { + *slot = middle_last_field(&values); + } + } + #[kernel] pub fn test_array_subslice_value(input: &[[u32; 4]], mut out: DisjointSlice) { let idx = thread::index_1d(); @@ -124,6 +253,60 @@ mod kernels { } } + #[kernel] + pub fn test_slice_subslice_from_end_mut(input: &[[u32; 4]], mut out: DisjointSlice) { + let idx = thread::index_1d(); + let i = idx.get(); + if i >= input.len() { + return; + } + + let mut values = input[i]; + let before = values; + + // Exercise two distinct from-end ConstantIndex offsets on the rebuilt + // middle subslice: offset 1 selects original element 2, offset 2 selects + // original element 1. + bump_slice_middle_last(&mut values); + bump_slice_middle_second_last(&mut values); + + if let Some(slot) = out.get_mut(idx) { + let d0 = values[0] - before[0]; + let d1 = values[1] - before[1]; + let d2 = values[2] - before[2]; + let d3 = values[3] - before[3]; + + // Pack every element's delta so collateral writes also fail the test. + *slot = d0 | (d1 << 8) | (d2 << 16) | (d3 << 24); + } + } + + #[kernel] + pub fn test_slice_subslice_from_end_runtime_len( + input: &[[u32; 4]], + mut out: DisjointSlice, + ) { + let idx = thread::index_1d(); + let i = idx.get(); + if i >= input.len() { + return; + } + + // The slice length is only known at run time, so the from-end index + // cannot fold to a constant element position. The kernel checks the + // buffer bound itself; the host oracle has no such check and panics + // on a bad length instead. + let row = input[i]; + let n = runtime_len(row); + if let Some(slot) = out.get_mut(idx) { + *slot = if n <= BUFFER_LEN { + from_end_at_runtime_len(row, n) + } else { + u32::MAX + }; + } + } + #[kernel] pub fn test_slice_subslice_mut(input: &[[u32; 4]], mut out: DisjointSlice) { let idx = thread::index_1d(); @@ -143,15 +326,17 @@ mod kernels { } } -const N: usize = 4; +/// Six lanes: column 0 of the rows runs from 10 to 15, so the runtime-length +/// case sees every length from 3 to 8. +const N: usize = 6; + +fn input_row(i: usize) -> [u32; 4] { + let i = i as u32; + [10 + i, 20 + i, 30 + i, 40 + i] +} fn make_input() -> Vec<[u32; 4]> { - (0..N) - .map(|i| { - let i = i as u32; - [10 + i, 20 + i, 30 + i, 40 + i] - }) - .collect() + (0..N).map(input_row).collect() } fn run_and_report(name: &str, stream: &Arc, launch: F, expected: E) -> bool @@ -190,6 +375,16 @@ fn main() { let mut all_pass = true; + all_pass &= run_and_report( + "slice subslice from-end field", + &stream, + |s, cfg, i, o| { + // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + unsafe { module.test_slice_subslice_from_end_field(s, cfg, i, o) }.expect("launch") + }, + |i| 160 + 2 * i as u32, + ); + all_pass &= run_and_report( "array value", &stream, @@ -240,6 +435,30 @@ fn main() { |_| 9, ); + all_pass &= run_and_report( + "slice subslice from-end mutable", + &stream, + |s, cfg, i, o| { + // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + unsafe { module.test_slice_subslice_from_end_mut(s, cfg, i, o) }.expect("launch") + }, + |_| (13 << 8) | (11 << 16), + ); + + all_pass &= run_and_report( + "slice subslice from-end runtime", + &stream, + |s, cfg, i, o| { + // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. + unsafe { module.test_slice_subslice_from_end_runtime_len(s, cfg, i, o) } + .expect("launch") + }, + |i| { + let row = input_row(i); + from_end_at_runtime_len(row, runtime_len(row)) + }, + ); + if all_pass { println!("\nSUCCESS: all Subslice projection cases passed"); } else { 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 2ce3306137..1702573ea1 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/switchint_128bit/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/Cargo.toml index e5416bb132..fe40f1a447 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/switchint_128bit/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.lock index 3b56442154..95370fb238 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.toml index 5cbade255f..fc74a30cf3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/swizzle_smem/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime (with HMM support) -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 1f5d931f29..83d5191784 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.toml index c0804c04d7..a8181b5ed3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/Cargo.toml @@ -15,4 +15,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/README.md index e0964fc7f9..434636a626 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05/README.md @@ -12,11 +12,13 @@ Tests tcgen05 (Tensor Core Gen 5) infrastructure for Blackwell GPUs. This is the 2. **tcgen05_alloc_test**: TMEM (Tensor Memory) allocation/deallocation 3. **tcgen05_commit_test**: Commit with mbarrier integration 4. **tcgen05_mma_minimal**: Full MMA pipeline (alloc → copy → MMA → read → dealloc) +5. **tmem_guard_2d_test**: `TmemGuard` alloc/dealloc in a (32, 4, 1) block; the guard selects the warp by block-linear thread id, so one hardware warp allocates +6. **tmem_guard_2d_warp3_test**: the same block through `alloc_by(3)` / `dealloc_by(3)`; a separate kernel, since repeated TMEM allocations in one CTA must come from the same warp ### cta_group::2 (CTA pairs) -5. **tcgen05_alloc_cg2_test**: Cooperative TMEM alloc/dealloc across a 2-CTA cluster -6. **tcgen05_mma_cg2_test**: Cooperative MMA with multicast commit across the pair +7. **tcgen05_alloc_cg2_test**: Cooperative TMEM alloc/dealloc across a 2-CTA cluster +8. **tcgen05_mma_cg2_test**: Cooperative MMA with multicast commit across the pair CTA pairs place 2 CTAs on adjacent SMs (a TPC). They cooperate on larger MMA tiles — each SM's tensor core handles its half of the rows. All `tcgen05` @@ -55,7 +57,9 @@ let desc = Tcgen05SmemDescriptor::builder() // Address 0x0 is a valid TMEM base — don't treat it as failure. static mut TMEM_ADDR: SharedArray = SharedArray::UNINIT; -// Warp-synchronous: all 32 threads must execute together +// Warp-synchronous: all 32 threads must execute together. +// `warp_id` is `threadIdx.x / 32`, one hardware warp only in a 1D block; +// `TmemGuard` selects the warp by block-linear thread id for any block. if warp_id == 0 { tcgen05_alloc(&raw mut TMEM_ADDR as *mut u32, 64); // or for CTA pairs: tcgen05_alloc_cg2(..., 512); @@ -133,6 +137,11 @@ TMEM address: 0x00000000 --- Test: tcgen05 MMA Minimal --- ✓ MMA minimal test executed successfully +--- Test: TmemGuard in a (32, 4, 1) block --- +TMEM address (warp 0): 0x00000000 +TMEM address (warp 3): 0x00000000 +✓ TmemGuard allocated and freed from block-linear warps 0 and 3, one kernel each + --- CTA Pair (cta_group::2) Tests --- --- Test: tcgen05 TMEM Alloc cta_group::2 --- 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 824b708892..acb06ddc36 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 @@ -14,6 +14,9 @@ //! - Tcgen05SmemDescriptor / Tcgen05InstructionDescriptor builders //! - tcgen05_mma_ws_f16() //! - tcgen05_cp_smem_to_tmem() / tcgen05_ld_16x256b_pure() +//! - TmemGuard::alloc() / dealloc() and alloc_by(3) / dealloc_by(3) in a +//! (32, 4, 1) block, one TMEM allocation per kernel, selecting the warp by +//! block-linear thread id //! //! **cta_group::2 (CTA pairs):** //! - tcgen05_alloc_cg2() / tcgen05_dealloc_cg2() @@ -35,10 +38,11 @@ use cuda_device::barrier::Barrier; use cuda_device::shared::{SharedArray, cvta_generic_to_shared_offset}; use cuda_device::tcgen05::{ self, Tcgen05AccumulatorType, Tcgen05ElementType, Tcgen05InstructionDescriptor, - Tcgen05MmaShape, Tcgen05SmemDescriptor, Tcgen05SwizzleMode, tcgen05_alloc, tcgen05_alloc_cg2, - tcgen05_commit, tcgen05_commit_multicast_cg2, tcgen05_cp_smem_to_tmem, tcgen05_dealloc, - tcgen05_dealloc_cg2, tcgen05_fence_after_thread_sync, tcgen05_fence_before_thread_sync, - tcgen05_ld_16x256b_pure, tcgen05_load_wait, tcgen05_mma_f16_cg2, tcgen05_mma_ws_f16, + Tcgen05MmaShape, Tcgen05SmemDescriptor, Tcgen05SwizzleMode, TmemGuard, TmemUninit, + tcgen05_alloc, tcgen05_alloc_cg2, tcgen05_commit, tcgen05_commit_multicast_cg2, + tcgen05_cp_smem_to_tmem, tcgen05_dealloc, tcgen05_dealloc_cg2, tcgen05_fence_after_thread_sync, + tcgen05_fence_before_thread_sync, tcgen05_ld_16x256b_pure, tcgen05_load_wait, + tcgen05_mma_f16_cg2, tcgen05_mma_ws_f16, }; use cuda_device::{CuSimd, DisjointSlice, cluster, cluster_launch, kernel, thread, warp}; use cuda_host::cuda_module; @@ -119,6 +123,67 @@ mod kernels { } } + /// `TmemGuard::alloc()` and `dealloc()` in a (32, 4, 1) block. + /// + /// Every thread calls both. The guard selects the allocating warp by + /// block-linear thread id, so only warp 0, the y == 0 row, runs + /// `tcgen05.alloc` and `tcgen05.dealloc`; under a `threadIdx.x / 32` + /// selection all four rows would allocate. Thread (0, 0, 0) stores the + /// TMEM address to `output[0]`. + /// + /// The kernel allocates once. The CUTLASS TMEM allocator requires + /// repeated allocations in one CTA to come from the same warp, so the + /// warp 3 case lives in its own kernel, `tmem_guard_2d_warp3_test`. + #[kernel] + pub unsafe fn tmem_guard_2d_test(mut output: DisjointSlice) { + static mut TMEM_SLOT: SharedArray = SharedArray::UNINIT; + + unsafe { + let first = thread::threadIdx_x() == 0 && thread::threadIdx_y() == 0; + let slot = &raw mut TMEM_SLOT as *mut u32; + + if first { + *slot = 0xDEADBEEF; + } + thread::sync_threads(); + + // Block-linear warp 0 allocates and frees. + let tmem = TmemGuard::::from_static(slot).alloc(); + if first { + *output.get_unchecked_mut(0) = tmem.raw_address(); + } + let _done = tmem.dealloc(); + } + } + + /// `TmemGuard::alloc_by(3)` and `dealloc_by(3)` in a (32, 4, 1) block. + /// + /// The same block as `tmem_guard_2d_test`, with the job handed to + /// block-linear warp 3, the y == 3 row. Thread (0, 0, 0) stores the TMEM + /// address to `output[1]`, next to the warp 0 result. One allocation per + /// kernel, for the reason given on `tmem_guard_2d_test`. + #[kernel] + pub unsafe fn tmem_guard_2d_warp3_test(mut output: DisjointSlice) { + static mut TMEM_SLOT: SharedArray = SharedArray::UNINIT; + + unsafe { + let first = thread::threadIdx_x() == 0 && thread::threadIdx_y() == 0; + let slot = &raw mut TMEM_SLOT as *mut u32; + + if first { + *slot = 0xDEADBEEF; + } + thread::sync_threads(); + + // Block-linear warp 3 allocates and frees. + let tmem = TmemGuard::::from_static(slot).alloc_by(3); + if first { + *output.get_unchecked_mut(1) = tmem.raw_address(); + } + let _done = tmem.dealloc_by(3); + } + } + /// Test kernel for tcgen05 commit with mbarrier. #[kernel] pub unsafe fn tcgen05_commit_test(mut output: DisjointSlice) { @@ -867,6 +932,7 @@ fn main() -> Result<(), Box> { run_tcgen05_alloc_test(&stream, &module)?; run_tcgen05_commit_test(&stream, &module)?; run_tcgen05_mma_minimal_test(&stream, &module)?; + run_tmem_guard_2d_test(&stream, &module)?; println!("\n--- CTA Pair (cta_group::2) Tests ---\n"); run_tcgen05_alloc_cg2_test(&stream, &module)?; @@ -999,6 +1065,43 @@ fn run_tcgen05_alloc_test( Ok(()) } +fn run_tmem_guard_2d_test( + stream: &Arc, + module: &kernels::LoadedModule, +) -> Result<(), Box> { + println!("--- Test: TmemGuard in a (32, 4, 1) block ---\n"); + + let mut output = DeviceBuffer::::zeroed(stream, 2)?; + let cfg = LaunchConfig { + block_dim: (32, 4, 1), + grid_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + + // One TMEM allocation per kernel: warp 0 in the first launch, warp 3 in + // the second. Both kernels write into `output`, at index 0 and index 1. + println!("Launching tmem_guard_2d_test kernel..."); + unsafe { module.tmem_guard_2d_test((stream).as_ref(), cfg, &mut output) }?; + println!("Launching tmem_guard_2d_warp3_test kernel..."); + unsafe { module.tmem_guard_2d_warp3_test((stream).as_ref(), cfg, &mut output) }?; + + stream.synchronize()?; + + let host_output = output.to_host_vec(stream)?; + println!("TMEM address (warp 0): 0x{:08x}", host_output[0]); + println!("TMEM address (warp 3): 0x{:08x}", host_output[1]); + + // A slot that still holds the sentinel means the selected warp never ran + // tcgen05.alloc. This is the 2D block the guard's warp selection exists for, + // so it is a failure, not a note. + if host_output.contains(&0xDEADBEEF) { + return Err("TmemGuard left a TMEM slot holding the sentinel in a (32, 4, 1) block".into()); + } + println!("✓ TmemGuard allocated and freed from block-linear warps 0 and 3, one kernel each\n"); + + Ok(()) +} + fn run_tcgen05_commit_test( stream: &Arc, module: &kernels::LoadedModule, 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 8c2269f5d1..b25fb143ad 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tcgen05_matmul/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/Cargo.toml index a51f7be491..c13a3452e2 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tcgen05_matmul/Cargo.toml @@ -15,5 +15,5 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } half = "2.4" 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 34faff57b2..608a68f483 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/thread_runs/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/thread_runs/Cargo.toml index f255aec5a8..ffac18a508 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/thread_runs/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/thread_runs/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.lock index 4ba660c11e..1686557131 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.toml index d1aa6d24fe..580de65dd1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tiled_gemm/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 new file mode 100644 index 0000000000..7696b70a7f --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/Cargo.lock @@ -0,0 +1,735 @@ +# 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 = "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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[package]] +name = "cuda-macros" +version = "0.2.1" +dependencies = [ + "oxide-tile-export", + "proc-macro2", + "quote", + "reserved-oxide-symbols", + "syn", +] + +[[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 = [ + "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 = "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 = "tilelib_relu" +version = "0.1.0" +dependencies = [ + "cuda-artifact-finalizer", + "cuda-core", + "cuda-device", + "oxide-tile-export", +] + +[[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/crates/rustc-codegen-cuda/examples/tilelib_relu/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/Cargo.toml new file mode 100644 index 0000000000..4ca1afc7b9 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "tilelib_relu" +version = "0.1.0" +edition = "2024" +authors = ["Nihal Pasham "] +license = "Apache-2.0" + +# Mark as standalone crate (not part of parent workspace) +[workspace] + +# A tile library written in Rust: #[tile_export] functions a cuTile program +# calls through `foreign_call`. The host program proves the native SIMT ABI on +# this machine's GPU through a generated CUDA C++ caller: +# cargo oxide run tilelib_relu --emit-nvvm-ir --arch sm_120 + +[dependencies] +# Device-side intrinsics (#[tile_export], warp shuffles, etc.). +cuda-device = { path = "../../../cuda-device" } + +# The descriptions the backend writes next to the device artifact. The host +# derives the C++ prototypes from them instead of restating the ABI by hand. +oxide-tile-export = { path = "../../../oxide-tile-export" } + +# libNVVM and nvJitLink without the CUDA driver: NVVM IR to LTOIR, then the +# two-input LTOIR link into a cubin, the same calls cuda-host makes. +cuda-artifact-finalizer = { path = "../../../cuda-artifact-finalizer" } + +# Host-side CUDA runtime: context, device buffers, raw kernel launches. +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/README.md new file mode 100644 index 0000000000..95bff63d19 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/README.md @@ -0,0 +1,136 @@ +# tilelib_relu + +A tile library written in Rust. Three `#[tile_export]` functions compile to +one LTOIR module on the native SIMT ABI, and `tilelib_relu.tileexports.json` +describes them. A Tile kernel reaches them once the two are packaged as +`tilelib_relu.tilelib`, which needs a CUDA 13.5 or newer tileiras: from Tile +C++ through a foreign declaration (`cpp/tile_caller.cu`), or from cuTile +Python through `foreign_call` (`python/call_from_cutile.py`). The host +program here proves the ABI on any CUDA 13.4 machine, through a generated +CUDA C++ caller that declares the implementations the way the Tile compiler +would. + +```text +src/main.rs: #[tile_export] relu, scale_in_place, warp_sum + | + | cargo oxide run tilelib_relu --emit-nvvm-ir --arch sm_120 + v +tilelib_relu.ll ---libNVVM---> tilelib_relu.ltoir tilelib_relu.tileexports.json + (NVVM IR, SIMT ABI) (device code) (one description per export) + | | + +---------------------------+------------------------------+ + | | + v v +consumer 1: this host program (CUDA 13.4 is enough) consumer 2: cuTile (CUDA 13.5+) + descriptions -> mock/caller.cu cargo oxide tilelib tilelib_relu + nvcc -dc --keep -gencode arch=compute_120,code=lto_120 -> tilelib_relu.tilelib + nvJitLink(tilelib_relu.ltoir + caller.ltoir) -> cubin python/call_from_cutile.py + launch call_, compare, PASS x3 tp.foreign_call(symbol_name="relu") +``` + +## The exports + +| cuTile name | Rust | implementation the call resolves to | +|------------------|-----------------------------------------|----------------------------------------------------| +| `relu` | `fn relu(x: f32) -> f32` | `void relu_impl(float x, float *relu_out)` | +| `scale_in_place` | `fn scale(x: &mut f32, factor: f32)` | `void scale_in_place_impl(float *x, float factor)` | +| `warp_sum` | `fn warp_sum(x: f32) -> f32`, warp sync | `void warp_sum_impl(float x, float *warp_sum_out)` | + +`in` scalars travel by value, `&mut` parameters as pointers, and a return +value becomes a trailing `out` pointer. The host program does not restate +that table: it derives every prototype in `mock/caller.cu` from the +descriptions through `oxide_tile_export::TileExportDesc::abi_params`, so the +generated file is the ABI as the tooling understands it. It is written when +the program runs and is not tracked, like the other generated artifacts. + +The generated kernels launch one thread per invocation, tile size divided by +the export's footprint, and pass each thread the first of its footprint +elements; the three exports here have footprint 1. An export with `f16` +elements is skipped with a `SKIP` line and counted in the `SUCCESS` line +rather than failed, because nvcc has no scalar type that lowers to LLVM +`half`. + +## Run + +```bash +# 1. Build the library and prove the ABI on the GPU (needs nvcc, libNVVM, +# nvJitLink and a GPU matching --arch; CUDA 13.4 is enough). +cargo oxide run tilelib_relu --emit-nvvm-ir --arch sm_120 + +# 2. Package it as a tile library (needs a tileiras that lists 13.5 or newer). +cargo oxide tilelib tilelib_relu --arch sm_120 + +# 3. Call it from a Tile C++ kernel (needs the same toolkit's nvcc, which +# knows --enable-tile; the kernel runs on the GPU and checks all three exports). +nvcc --enable-tile --tile-preview simt -Xtileiras=-l=tilelib_relu.tilelib \ + -std=c++20 -arch=sm_120 -o tile_caller cpp/tile_caller.cu && ./tile_caller + +# 4. Or call it from cuTile Python (checked against cuTile Python 1.7). +python python/call_from_cutile.py --dry-run # bytecode only, no tileiras +python python/call_from_cutile.py # full run, needs step 2 +``` + +Step 1 is not a Tile kernel. The host program stands in for the Tile +compiler: it writes a CUDA C++ caller from the descriptions and links it +against the LTOIR, which proves the ABI on CUDA 13.4. The end-to-end test, a +Tile kernel calling the Rust functions, is step 2 followed by step 3 or 4. +Run step 2 after step 1: `cargo oxide run` and `cargo oxide build` clear the +example's generated artifacts, the `.tilelib` among them. When the toolkit +that brings the 13.5 assembler is not the one on `PATH`, point +`CUTILE_TILEIRAS_PATH` at its `tileiras` for step 2 and run its `nvcc` for +step 3. + +Step 3 prints one `PASS` line per export and then +`SUCCESS: 3/3 tile exports called from a Tile C++ kernel`. Each foreign +declaration in `cpp/tile_caller.cu` mirrors a declaration of the library: +parameters the library reads carry no attribute, parameters it writes carry +`[[cutile::experimental_inout]]` and come back as the result, the way the +MathDx tile libraries declare theirs. + +Step 1 prints one line per export and then the marker: + +```text +PASS relu: 128 elements through `relu_impl`, relu of a ramp from -4 to 4 +PASS scale_in_place: 128 elements through `scale_in_place_impl`, the ramp scaled in place by 0.5 +PASS warp_sum: 128 elements through `warp_sum_impl`, 32-lane sums over 0..128, every lane holding its warp's total +SUCCESS: 3/3 tile exports verified through the native ABI +``` + +Without nvcc it prints `skipping: nvcc not found` and exits 0. + +## What PASS proves + +- The wrapper `#[tile_export]` generates has the calling convention the Tile + compiler derives from the same description: a C++ caller written from the + description alone links against it and gets the right numbers. +- The link works the way a tile library is consumed: two LTOIR inputs, + symbols resolved by nvJitLink, no PTX in between. +- `sync = warp` behaves when every lane of a warp calls together: all 32 + lanes receive the warp total. + +All three checks compare bit for bit; the inputs are chosen so f32 cannot +round (relu copies or zeroes, 0.5 halves, the warp sums are small integers). + +## What still needs CUDA 13.5 or newer + +- `tilelib_relu.tilelib` itself: `tileiras --preview-create-tile-lib` exists + from 13.5 on, and an assembler reads preview content only at the version + it calls current, so the descriptions are written at the resolved + assembler's newest listed version. When the installed tileiras lists only + 13.4, + `cargo oxide tilelib` stops after writing the descriptions, says so, and + exits 2; with `--library-optional` the descriptions alone count as + success and it exits 0, which is how the smoketest runs it. +- The Tile C++ caller: nvcc gains `--enable-tile` and `--tile-preview simt` + with the same toolkit that brings the assembler, and the program and the + library then agree on the bytecode version by construction. +- The cuTile side: `foreign_call` is a 13.5 bytecode feature. The dry run + compiles it on 13.4 because it never calls tileiras. It was checked + against cuTile Python 1.7 (checkout 770b493) with the `cuda-tile-preview` + package built from that checkout's `experimental/tile_preview` folder: + `cuda.tile_preview` is not on PyPI, and `call_from_cutile.py` imports + private `cuda.tile._*` modules (`_bytecode`, `_cext`, `_compile`) that can + move between releases. The full run also needs a cuTile Python that lists + the assembler's current bytecode version: the script lists every version + its encoder knows, and the probe against tileiras picks the newest one it + accepts. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/cpp/tile_caller.cu b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/cpp/tile_caller.cu new file mode 100644 index 0000000000..31e4942713 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/cpp/tile_caller.cu @@ -0,0 +1,123 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +// A Tile C++ kernel that calls the three Rust exports of tilelib_relu.tilelib. +// +// nvcc --enable-tile --tile-preview simt -Xtileiras=-l=tilelib_relu.tilelib \ +// -std=c++20 -arch=sm_120 -o tile_caller cpp/tile_caller.cu +// +// Needs a toolkit whose nvcc knows `--enable-tile` and whose tileiras built +// the library (CUDA 13.5 or newer). Each foreign declaration mirrors one +// declaration in the library's interface: parameters the library reads carry +// no attribute, parameters it writes carry `[[cutile::experimental_inout]]` +// and come back as the result. The kernel runs one 128-element tile through +// each export; the host checks every element bit for bit. + +#include +#include + +#include +#include +#include + +namespace ct = cuda::tiles; + +inline constexpr int tile_size = 128; +using tile_f32 = ct::tile>; + +// @relu(%x: tile<128xf32> in, %relu_out: tile<128xf32> out) -> tile<128xf32> +__experimental_foreign_decl__ __experimental_foreign_name__("relu") + tile_f32 relu(tile_f32 x, [[cutile::experimental_inout]] tile_f32 relu_out); + +// @scale_in_place(%x: tile<128xf32> inout, %factor: tile<128xf32> in) -> tile<128xf32> +__experimental_foreign_decl__ __experimental_foreign_name__("scale_in_place") + tile_f32 scale_in_place([[cutile::experimental_inout]] tile_f32 x, tile_f32 factor); + +// @warp_sum(%x: tile<128xf32> in, %warp_sum_out: tile<128xf32> out) -> tile<128xf32>, +// implemented with synchronization: 32 lanes call the Rust function together. +__experimental_foreign_decl__ __experimental_foreign_name__("warp_sum") + tile_f32 warp_sum(tile_f32 x, [[cutile::experimental_inout]] tile_f32 warp_sum_out); + +[[cutile::hint(0, num_worker_warps_per_cta = 4)]] __tile_global__ void call_exports( + float* ramp, float* factor, float* lanes, float* relu_out, float* scale_out, float* warp_out) { + auto view = [](float* data) { + return ct::partition_view { ct::tensor_span { data, ct::shape {} }, ct::shape {} }; + }; + auto relu_view = view(relu_out); + auto scale_view = view(scale_out); + auto warp_view = view(warp_out); + + relu_view.store(relu(view(ramp).load(0), relu_view.load(0)), 0); + scale_view.store(scale_in_place(view(ramp).load(0), view(factor).load(0)), 0); + warp_view.store(warp_sum(view(lanes).load(0), warp_view.load(0)), 0); +} + +#define CHECK_CUDA(expression) \ + do { \ + const cudaError_t cuda_status = (expression); \ + if (cuda_status != cudaSuccess) { \ + std::fprintf(stderr, "%s failed: %s\n", #expression, cudaGetErrorString(cuda_status)); \ + return EXIT_FAILURE; \ + } \ + } while (0) + +int main() { + std::array ramp {}; + std::array factor {}; + std::array lanes {}; + for (int i = 0; i < tile_size; ++i) { + ramp[i] = -4.0f + 8.0f * float(i) / float(tile_size - 1); + factor[i] = 0.5f; + lanes[i] = float(i); + } + + float* device[6] = {}; + for (float*& pointer : device) CHECK_CUDA(cudaMalloc(&pointer, sizeof(ramp))); + CHECK_CUDA(cudaMemcpy(device[0], ramp.data(), sizeof(ramp), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(device[1], factor.data(), sizeof(factor), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(device[2], lanes.data(), sizeof(lanes), cudaMemcpyHostToDevice)); + for (int i = 3; i < 6; ++i) CHECK_CUDA(cudaMemset(device[i], 0, sizeof(ramp))); + + call_exports<<<1>>>(device[0], device[1], device[2], device[3], device[4], device[5]); + CHECK_CUDA(cudaGetLastError()); + CHECK_CUDA(cudaDeviceSynchronize()); + + std::array relu_out {}; + std::array scale_out {}; + std::array warp_out {}; + CHECK_CUDA(cudaMemcpy(relu_out.data(), device[3], sizeof(relu_out), cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy(scale_out.data(), device[4], sizeof(scale_out), cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy(warp_out.data(), device[5], sizeof(warp_out), cudaMemcpyDeviceToHost)); + + int failures = 0; + for (int i = 0; i < tile_size; ++i) { + const float expected_relu = ramp[i] > 0.0f ? ramp[i] : 0.0f; + const float expected_scale = ramp[i] * 0.5f; + // Lane values 32w .. 32w+31 sum to 32 * 32w + 496. + const float expected_warp = float(32 * 32 * (i / 32) + 496); + if (relu_out[i] != expected_relu) { + std::fprintf(stderr, "relu mismatch at %d: got %f expected %f\n", i, relu_out[i], expected_relu); + ++failures; + } + if (scale_out[i] != expected_scale) { + std::fprintf(stderr, "scale_in_place mismatch at %d: got %f expected %f\n", i, scale_out[i], expected_scale); + ++failures; + } + if (warp_out[i] != expected_warp) { + std::fprintf(stderr, "warp_sum mismatch at %d: got %f expected %f\n", i, warp_out[i], expected_warp); + ++failures; + } + } + for (float* pointer : device) CHECK_CUDA(cudaFree(pointer)); + if (failures != 0) { + std::fprintf(stderr, "FAIL: %d mismatches\n", failures); + return EXIT_FAILURE; + } + std::printf("PASS relu: 128 elements through the tile library\n"); + std::printf("PASS scale_in_place: 128 elements through the tile library\n"); + std::printf("PASS warp_sum: 128 elements through the tile library, every lane holding its warp's total\n"); + std::printf("SUCCESS: 3/3 tile exports called from a Tile C++ kernel\n"); + return EXIT_SUCCESS; +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/python/call_from_cutile.py b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/python/call_from_cutile.py new file mode 100644 index 0000000000..7a9b6e4beb --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/python/call_from_cutile.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Call tilelib_relu's `relu` from a cuTile kernel through `foreign_call`. + + python call_from_cutile.py --dry-run # bytecode only; runs on CUDA 13.4 + python call_from_cutile.py [path/to.tilelib] # full run; needs the CUDA 13.5 tileiras + +The kernel loads one 128-float tile, hands it to the library's `relu` +declaration together with a zero tile for the result slot, and stores what +comes back. `--dry-run` compiles that kernel to Tile IR bytecode without +tileiras (`return_cubin=False`) and checks that the program really calls +`relu`: the caller half of the contract, produced on this machine. The full +run also needs the library half, `tilelib_relu.tilelib`, which only a CUDA +13.5 tileiras can build (`cargo oxide tilelib tilelib_relu --arch sm_120`), +and a cuTile Python whose supported bytecode versions include 13.5. + +Checked against cuTile Python 1.7 (checkout 770b493) with the +`cuda-tile-preview` package from that checkout's experimental/tile_preview +folder; `cuda.tile_preview` is not on PyPI. The dry run reaches into private +`cuda.tile._*` modules (`_bytecode`, `_cext`, `_compile`), which can move +between releases. +""" + +from __future__ import annotations + +import argparse +import os +import pathlib +import shutil +import subprocess +import sys +import tempfile + +import numpy as np + +import cuda.tile as ct +from cuda import tile_preview as tp + +TILE = 128 +LIBRARY_VERSION = "0.1.0" +HERE = pathlib.Path(__file__).resolve().parent +DEFAULT_TILELIB = HERE.parent / "tilelib_relu.tilelib" +# preview$call in the bytecode opcode table; upstream's test_foreign_call.py +# looks for the same varint-encoded value. +PREVIEW_CALL_OPCODE = 0x4002 + +# Set by main() before the kernel is traced. foreign_call reads the library +# as a compile-time constant, so it has to be a module global, not an argument. +PROVIDER: tp.Tilelib | None = None + + +def relu_tile(x, y): + """One block runs `relu` over one 128-float tile of `x` into `y`.""" + pid = ct.bid(0) + tile = ct.load(x, index=(pid,), shape=(TILE,)) + # Every declaration argument needs an input tile, the `out` slot included. + result_slot = ct.zeros((TILE,), dtype=ct.float32) + (out,) = tp.foreign_call( + tilelibs=(PROVIDER,), + symbol_name="relu", + constant_params=(), + inputs=(tile, result_slot), + output_types=((ct.float32, (TILE,)),), + ) + ct.store(y, index=(pid,), tile=out) + + +def dry_run() -> None: + """Compile the kernel to 13.5 bytecode and show the preview$call to `relu` is in it.""" + from cuda.tile._bytecode.basic import encode_varint + from cuda.tile._bytecode.version import BytecodeVersion + from cuda.tile._cext import CallingConvention + from cuda.tile._compile import compile_tile + from cuda.tile.compilation import ArrayConstraint, KernelSignature + + constraint = ArrayConstraint( + dtype=ct.float32, + ndim=1, + index_dtype=ct.int32, + stride_lower_bound_incl=0, + alias_groups=(), + may_alias_internally=False, + ) + signature = KernelSignature( + [constraint, constraint], CallingConvention.cutile_python_v2(), symbol="relu_tile" + ) + result = compile_tile( + relu_tile, + [signature], + bytecode_version=BytecodeVersion.V_13_5, + return_final_ir=True, + return_bytecode=True, + return_cubin=False, + ) + bytecode = bytes(result.bytecode) + major, minor = bytecode[8], bytecode[9] + preview = int.from_bytes(bytecode[10:12], "little") & 1 + opcode = bytearray() + encode_varint(PREVIEW_CALL_OPCODE, opcode) + has_call_op = bytes(opcode) in bytecode + # The callee's name is read from the final IR, not from the bytecode's + # string table: the kernel symbol `relu_tile` contains `relu`, so a bare + # search of the bytes would pass without any call. + ir_text = "\n".join(str(block) for block in result.final_ir) + calls_relu = "preview_foreign_call" in ir_text and 'symbol_name="relu"' in ir_text + print(f"bytecode: {len(bytecode)} bytes, header {bytecode[8:12].hex()} " + f"(version {major}.{minor}, preview flag {'set' if preview else 'clear'}), " + f"preview$call opcode {'present' if has_call_op else 'absent'}; " + f"final IR {'calls' if calls_relu else 'does not call'} `relu`") + if (major, minor) != (13, 5) or not preview or not has_call_op or not calls_relu: + print("FAIL: the program does not carry a 13.5 preview call to `relu`") + sys.exit(1) + print("PASS: the caller side compiles; the program calls `relu` through " + "preview$call (dry run, no cubin)") + + +def full_run() -> None: + """Link the library with tileiras, run on the GPU, compare with numpy.""" + import cupy as cp + from cuda.tile import _compile as compile_module + from cuda.tile._bytecode.version import BytecodeVersion + + # A launched kernel is compiled at the newest bytecode version this cuTile + # Python lists and the installed tileiras accepts. The 1.7 checkout lists + # versions above 13.4 only when its extension is built with development + # features, so every version its encoder knows is listed here; the probe + # against tileiras still decides which one is used. foreign_call needs + # 13.5, and an assembler reads preview content only at the version it + # calls current, so this cuTile Python must encode that version too. A + # release that lists these versions itself needs no such lines. + for version in BytecodeVersion: + if (version > BytecodeVersion.V_13_4 + and version not in compile_module._SUPPORTED_VERSIONS): + compile_module._SUPPORTED_VERSIONS.append(version) + + relu_kernel = ct.kernel(relu_tile) + n = 32 * TILE + x = cp.linspace(-4.0, 4.0, n, dtype=cp.float32) + y = cp.zeros_like(x) + try: + ct.launch(cp.cuda.get_current_stream(), (n // TILE, 1, 1), relu_kernel, (x, y)) + except Exception as error: # noqa: BLE001 (every failure here is reported, not hidden) + explain(error) + sys.exit(1) + expected = np.maximum(cp.asnumpy(x), 0.0) + got = cp.asnumpy(y) + np.testing.assert_array_equal(got, expected) + print(f"PASS: relu from the tile library matches numpy on {n} elements") + + +def explain(error: Exception) -> None: + """Report the failure and, when it is the toolkit, say so plainly.""" + print(f"FAIL: {type(error).__name__}: {error}") + versions = tileiras_versions() + if versions and "13.5" not in versions: + print(f"The installed tileiras lists bytecode versions {', '.join(versions)}. " + "foreign_call needs 13.5, which ships with CUDA 13.5 or newer; until then " + "only --dry-run can run on this machine.") + else: + print("A 13.5 tileiras is also only used when this cuTile Python lists 13.5 " + "among its supported bytecode versions.") + + +def tileiras_versions() -> list[str]: + """The bytecode versions the tileiras on this machine accepts, or [].""" + candidates = [shutil.which("tileiras")] + for root in (os.environ.get("CUDA_HOME"), os.environ.get("CUDA_PATH"), "/usr/local/cuda"): + if root: + candidates.append(os.path.join(root, "bin", "tileiras")) + for path in candidates: + if not path or not os.path.isfile(path): + continue + try: + output = subprocess.run([path, "--list-versions"], capture_output=True, + text=True, check=True, timeout=30).stdout + except (OSError, subprocess.SubprocessError): + continue + return [line.strip() for line in output.splitlines() if line.strip()] + return [] + + +def main() -> None: + global PROVIDER + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("tilelib", nargs="?", default=DEFAULT_TILELIB, type=pathlib.Path, + help="the tile library (default: ../tilelib_relu.tilelib)") + parser.add_argument("--dry-run", action="store_true", + help="compile to bytecode only; no tileiras, no GPU") + args = parser.parse_args() + + if args.dry_run and not args.tilelib.is_file(): + # Tilelib() insists that the file exists; the bytecode only records + # its path, so a placeholder stands in for the real library. + with tempfile.TemporaryDirectory() as scratch: + placeholder = pathlib.Path(scratch) / args.tilelib.name + placeholder.write_bytes(b"placeholder") + print(f"{args.tilelib} is not built yet; using a placeholder for the dry run") + PROVIDER = tp.Tilelib(path=placeholder, version=LIBRARY_VERSION) + dry_run() + return + + if not args.tilelib.is_file(): + print(f"FAIL: {args.tilelib} does not exist; build it with " + "`cargo oxide tilelib tilelib_relu --arch sm_120` (CUDA 13.5 tileiras)") + sys.exit(1) + PROVIDER = tp.Tilelib(path=args.tilelib, version=LIBRARY_VERSION) + if args.dry_run: + dry_run() + else: + full_run() + + +if __name__ == "__main__": + main() diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/src/main.rs new file mode 100644 index 0000000000..c6168616d6 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/src/main.rs @@ -0,0 +1,593 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! tilelib_relu: a tile library written in Rust, checked on the GPU. +//! +//! The device half is the three `#[tile_export]` functions below. Each one +//! becomes an implementation a cuTile program calls through `foreign_call`, +//! on the native SIMT ABI: `in` scalars by value, everything else behind a +//! pointer, the return value as a trailing `out` pointer. +//! +//! The host half proves that ABI without a CUDA 13.5 or newer tileiras. It reads the +//! descriptions the backend wrote, generates `mock/caller.cu` with the +//! prototypes the Tile compiler derives, compiles it with nvcc, links it with +//! the Rust LTOIR through nvJitLink, and checks the numbers on the GPU: +//! +//! cargo oxide run tilelib_relu --emit-nvvm-ir --arch sm_120 + +use cuda_artifact_finalizer::{ + CudaArch, DebugPolicy, FinalizationOptions, FinalizerOutput, LtoLinker, NamedInput, + NvvmCompiler, +}; +use cuda_core::embedded::{ArtifactCompileOptions, ArtifactDebugPolicy}; +use cuda_core::{CudaContext, CudaModule, DeviceBuffer, launch_kernel_on_stream}; +use cuda_device::{tile_export, warp}; +use oxide_tile_export::{ + AbiParamKind, DESCRIPTIONS_SUFFIX, TileAccess, TileElem, TileExportDesc, descriptions_from_json, +}; +use std::error::Error; +use std::ffi::c_void; +use std::fmt::Write as _; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::Arc; + +// ============================================================================ +// Device half: the tile library +// ============================================================================ + +/// Elementwise `max(x, 0)`. +#[tile_export(tile_shape = 128)] +pub fn relu(x: f32) -> f32 { + if x > 0.0 { x } else { 0.0 } +} + +/// Scale a tile in place. The cuTile side knows it as `scale_in_place`. +#[tile_export( + tile_shape = 128, + name = "scale_in_place", + symbol = "scale_in_place_impl" +)] +pub fn scale(x: &mut f32, factor: f32) { + *x *= factor; +} + +/// Sum over the 32 lanes of each warp; every lane gets the warp total. +#[tile_export(tile_shape = 128, sync = "warp")] +pub fn warp_sum(x: f32) -> f32 { + let mut total = x; + let mut offset = 16; + while offset > 0 { + total += warp::shuffle_xor_f32_sync(u32::MAX, total, offset); + offset /= 2; + } + total +} + +// ============================================================================ +// Host half: prove the ABI through a generated CUDA C++ caller +// ============================================================================ + +const CRATE: &str = "tilelib_relu"; + +type Fallible = Result>; + +/// How the run ended when nothing went wrong. +enum Outcome { + /// The exports that were called through the native ABI and checked, plus + /// the ones the generated caller cannot declare, each with its reason. + Verified { + verified: usize, + skipped: Vec<(String, String)>, + }, + /// The check cannot run on this machine; the reason is printed. + Skipped(String), +} + +fn main() { + match run() { + Ok(Outcome::Verified { verified, skipped }) => { + let total = verified + skipped.len(); + let mut line = + format!("SUCCESS: {verified}/{total} tile exports verified through the native ABI"); + if !skipped.is_empty() { + let names: Vec<&str> = skipped.iter().map(|(decl, _)| decl.as_str()).collect(); + let _ = write!(line, ", {} skipped ({})", skipped.len(), names.join(", ")); + } + println!("{line}"); + } + Ok(Outcome::Skipped(reason)) => println!("skipping: {reason}"), + Err(error) => { + println!("FAIL: {error}"); + std::process::exit(1); + } + } +} + +fn run() -> Fallible { + let dir = Path::new(env!("CARGO_MANIFEST_DIR")); + let (exports, skipped) = partition_callable(read_descriptions(dir)?); + if exports.is_empty() { + return Ok(Outcome::Skipped( + "every export uses half precision, which the generated caller cannot declare" + .to_string(), + )); + } + let options = read_build_policy(dir)?; + let Some(nvcc) = find_nvcc() else { + return Ok(Outcome::Skipped("nvcc not found".to_string())); + }; + + let caller_cu = dir.join("mock").join("caller.cu"); + write_if_changed(&caller_cu, &caller_source(&exports)?)?; + + let rust_ltoir = rust_ltoir(dir, &options)?; + let caller_ltoir = compile_caller(&nvcc, &caller_cu, options.target())?; + // Both halves are LTOIR, so nvJitLink resolves the implementation symbols + // at link time, the same way a real tile library is consumed. + let cubin = LtoLinker::discover()?.link_ltoir( + &[ + NamedInput::new("tilelib_relu.ltoir", &rust_ltoir), + NamedInput::new("caller.ltoir", &caller_ltoir), + ], + &options, + FinalizerOutput::Cubin, + )?; + println!( + "nvJitLink: {} bytes of cubin for {}", + cubin.len(), + options.target().sm() + ); + + let ctx = match CudaContext::new(0) { + Ok(ctx) => ctx, + Err(error) => return Ok(Outcome::Skipped(format!("no CUDA device: {error}"))), + }; + let (major, minor) = ctx.compute_capability()?; + let gpu = u32::try_from(major * 10 + minor)?; + if gpu != options.target().capability() { + return Ok(Outcome::Skipped(format!( + "the device code targets {}, this GPU is sm_{gpu}", + options.target().sm() + ))); + } + let module = ctx.load_module_from_image(&cubin)?; + for desc in &exports { + check_export(&ctx, &module, desc)?; + } + Ok(Outcome::Verified { + verified: exports.len(), + skipped, + }) +} + +/// Split the exports into the ones the generated caller can declare and the +/// ones it cannot, printing one `SKIP` line per export left out. nvcc has no +/// scalar type that lowers to LLVM `half` or `bfloat` (its `__half` is a +/// struct), so a half-precision export is reported and counted rather than +/// failed: the ABI defect that would be is in the mock, not in the export. +fn partition_callable( + exports: Vec, +) -> (Vec, Vec<(String, String)>) { + let mut callable = Vec::with_capacity(exports.len()); + let mut skipped = Vec::new(); + for desc in exports { + match desc + .params + .iter() + .find_map(|param| c_type(param.elem).err()) + { + Some(reason) => { + println!("SKIP {}: {reason}", desc.decl); + skipped.push((desc.decl, reason)); + } + None => callable.push(desc), + } + } + (callable, skipped) +} + +/// The descriptions the backend wrote next to the device artifact, validated +/// so the generated C++ can rely on the ABI rules. +fn read_descriptions(dir: &Path) -> Fallible> { + let path = dir.join(format!("{CRATE}.{DESCRIPTIONS_SUFFIX}")); + let text = std::fs::read_to_string(&path).map_err(|error| { + format!( + "could not read {} (build with `cargo oxide run {CRATE} --emit-nvvm-ir --arch sm_120`): {error}", + path.display() + ) + })?; + let exports = descriptions_from_json(&text)?; + if exports.is_empty() { + return Err(format!("{} describes no tile exports", path.display()).into()); + } + for desc in &exports { + desc.validate() + .map_err(|error| format!("description of `{}`: {error}", desc.decl))?; + } + Ok(exports) +} + +/// The target and compile policy recorded with the NVVM IR. The link must use +/// the same architecture and FMA setting as the build, or nvJitLink rejects +/// the input or changes the arithmetic. +fn read_build_policy(dir: &Path) -> Fallible { + let target_path = dir.join(format!("{CRATE}.target")); + let text = std::fs::read_to_string(&target_path) + .map_err(|error| format!("could not read {}: {error}", target_path.display()))?; + let target: CudaArch = text.lines().next().unwrap_or_default().trim().parse()?; + + let options_path = dir.join(format!("{CRATE}.options")); + let compile = match std::fs::read_to_string(&options_path) { + Ok(text) => ArtifactCompileOptions::from_sidecar_text(&text) + .map_err(|error| format!("invalid {}: {error}", options_path.display()))?, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => ArtifactCompileOptions::new(), + Err(error) => { + return Err(format!("could not read {}: {error}", options_path.display()).into()); + } + }; + let debug = match compile.debug_policy() { + ArtifactDebugPolicy::None => DebugPolicy::None, + ArtifactDebugPolicy::LineTables => DebugPolicy::LineTables, + ArtifactDebugPolicy::Full => DebugPolicy::Full, + }; + Ok(FinalizationOptions::new(target) + .with_fma_contraction(compile.fma_contraction_enabled()) + .with_debug_policy(debug)) +} + +/// The Rust half as LTOIR. `cargo oxide tilelib` and `emit-ltoir` leave a +/// `.ltoir` next to the `.ll`; a plain `--emit-nvvm-ir` build leaves only the +/// `.ll`, so it is compiled here with libNVVM, the call cargo oxide makes. A +/// `.ltoir` older than the `.ll` is stale and gets recompiled too. +fn rust_ltoir(dir: &Path, options: &FinalizationOptions) -> Fallible> { + let ll = dir.join(format!("{CRATE}.ll")); + let ltoir = dir.join(format!("{CRATE}.ltoir")); + if ltoir.is_file() && !older_than(<oir, &ll) { + println!("libNVVM: using {}", ltoir.display()); + return Ok(std::fs::read(<oir)?); + } + let ir = std::fs::read(&ll).map_err(|error| { + format!( + "could not read {} (build with `cargo oxide run {CRATE} --emit-nvvm-ir --arch sm_120`): {error}", + ll.display() + ) + })?; + let ltoir = + NvvmCompiler::discover()?.compile_nvvm_ir_to_ltoir(&format!("{CRATE}.ll"), &ir, options)?; + println!( + "libNVVM: {} compiled to LTOIR ({} bytes)", + ll.display(), + ltoir.len() + ); + Ok(ltoir) +} + +fn older_than(file: &Path, other: &Path) -> bool { + let modified = |path: &Path| std::fs::metadata(path).and_then(|meta| meta.modified()); + match (modified(file), modified(other)) { + (Ok(file_time), Ok(other_time)) => file_time < other_time, + _ => false, + } +} + +/// nvcc from `PATH`, then from the toolkit roots libNVVM is also looked up in. +fn find_nvcc() -> Option { + let on_path = std::env::var_os("PATH") + .map(|path| std::env::split_paths(&path).collect::>()) + .unwrap_or_default() + .into_iter() + .map(|dir| dir.join("nvcc")); + let in_toolkit = ["CUDA_TOOLKIT_PATH", "CUDA_HOME", "CUDA_PATH"] + .into_iter() + .filter_map(std::env::var_os) + .map(PathBuf::from) + .chain([PathBuf::from("/usr/local/cuda"), PathBuf::from("/opt/cuda")]) + .map(|root| root.join("bin").join("nvcc")); + on_path + .chain(in_toolkit) + .find(|candidate| candidate.is_file()) +} + +/// Generate `mock/caller.cu`: the implementation prototypes the Tile compiler +/// derives from the descriptions, and one kernel per export that calls the +/// implementation the way the compiler does over one tile. +fn caller_source(exports: &[TileExportDesc]) -> Fallible { + let mut out = String::new(); + out.push_str( + "// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n\ + // SPDX-License-Identifier: Apache-2.0\n\ + //\n\ + // Generated by the tilelib_relu host program from tilelib_relu.tileexports.json.\n\ + // Do not edit: it is rewritten whenever the descriptions change.\n\ + //\n\ + // Each kernel stands in for the Tile compiler: one thread per invocation of\n\ + // the implementation, launched as one block of tile-size / footprint threads.\n\ + // Scalar `in` parameters arrive by value, every other parameter as a pointer\n\ + // to the first of the invocation's footprint elements. A `sync = warp` export\n\ + // is called by every lane of each warp, which the full block provides.\n\n", + ); + for desc in exports { + // `validate` holds every parameter of an export to one footprint, so + // the first parameter's is the export's. + let footprint = desc.params[0].footprint; + let abi = desc.abi_params(); + let mut prototype = Vec::new(); + let mut kernel_params = Vec::new(); + let mut call_args = Vec::new(); + for (param, info) in abi.iter().zip(&desc.params) { + let c_type = c_type(param.kind.elem()).map_err(|reason| { + format!("`{}`: parameter `{}`: {reason}", desc.decl, param.name) + })?; + let constness = if info.access == TileAccess::In { + "const " + } else { + "" + }; + // Every parameter is a tile-size buffer to the kernel; by-value + // parameters are loaded from it, pointer parameters point into it. + kernel_params.push(format!("{constness}{c_type} *{}", param.name)); + match param.kind { + AbiParamKind::Scalar(_) => { + prototype.push(format!("{c_type} {}", param.name)); + call_args.push(format!("{}[i]", param.name)); + } + AbiParamKind::Pointer(_) => { + prototype.push(format!("{constness}{c_type} *{}", param.name)); + call_args.push(format!("&{}[i]", param.name)); + } + } + } + writeln!( + out, + "// {}: {} elements, footprint {footprint}, sync {}", + desc.decl, + desc.shape[0], + desc.sync.tile_ir_name() + )?; + writeln!( + out, + "extern \"C\" __device__ void {}({});", + desc.symbol, + prototype.join(", ") + )?; + writeln!( + out, + "extern \"C\" __global__ void call_{}({}) {{", + desc.decl, + kernel_params.join(", ") + )?; + if footprint == 1 { + writeln!( + out, + " const int i = blockIdx.x * blockDim.x + threadIdx.x;" + )?; + } else { + // Thread t owns elements [t * F, (t + 1) * F). + writeln!( + out, + " const int i = (blockIdx.x * blockDim.x + threadIdx.x) * {footprint};" + )?; + } + writeln!(out, " {}({});", desc.symbol, call_args.join(", "))?; + writeln!(out, "}}\n")?; + } + Ok(out) +} + +/// The C type of a tile element, or why nvcc has none: it has no scalar type +/// that lowers to LLVM `half` or `bfloat` (its `__half` is a struct), so +/// those two cannot be declared in the generated caller. +fn c_type(elem: TileElem) -> Result<&'static str, String> { + Ok(match elem { + TileElem::F32 => "float", + TileElem::F64 => "double", + TileElem::I8 => "signed char", + TileElem::I16 => "short", + TileElem::I32 => "int", + TileElem::I64 => "long long", + TileElem::U8 => "unsigned char", + TileElem::U16 => "unsigned short", + TileElem::U32 => "unsigned int", + TileElem::U64 => "unsigned long long", + TileElem::F16 | TileElem::BF16 => { + return Err(format!( + "{} has no by-value C type in nvcc; the mock covers f32, f64 and the integers", + elem.rust_name() + )); + } + }) +} + +/// Write `text` to `path` only when it differs, so a run that changes nothing +/// leaves the generated file alone. +fn write_if_changed(path: &Path, text: &str) -> Fallible<()> { + if std::fs::read_to_string(path).is_ok_and(|current| current == text) { + println!("caller: {} is current", path.display()); + return Ok(()); + } + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, text)?; + println!("caller: wrote {}", path.display()); + Ok(()) +} + +/// Compile the generated caller to LTOIR for `target` in a scratch directory, +/// removed on success and kept for inspection on failure. `-dc` leaves the +/// Rust symbols unresolved for the link; `--keep` makes nvcc leave the +/// `.ltoir` behind. +fn compile_caller(nvcc: &Path, caller_cu: &Path, target: &CudaArch) -> Fallible> { + let work = std::env::temp_dir().join(format!("{CRATE}-nvcc-{}", std::process::id())); + std::fs::create_dir_all(&work)?; + let sm = target.sm(); + let code = sm.trim_start_matches("sm_"); + + let mut command = Command::new(nvcc); + // The same host-compiler override the other LTOIR examples honor. + if let Some(ccbin) = std::env::var_os("NVCC_CCBIN").or_else(|| std::env::var_os("CUDAHOSTCXX")) + { + command.arg("-ccbin").arg(ccbin); + } + let output = command + .args(["-dc", "--keep", "--keep-dir"]) + .arg(&work) + .arg("-gencode") + .arg(format!("arch={},code=lto_{code}", target.compute())) + .arg(caller_cu) + .arg("-o") + .arg(work.join("caller.o")) + .current_dir(&work) + .output()?; + if !output.status.success() { + return Err(format!( + "nvcc failed ({}); intermediates kept in {}\n{}", + output.status, + work.display(), + String::from_utf8_lossy(&output.stderr) + ) + .into()); + } + let ltoir = std::fs::read(work.join("caller.ltoir"))?; + std::fs::remove_dir_all(&work)?; + println!("nvcc: caller.ltoir ({} bytes, lto_{code})", ltoir.len()); + Ok(ltoir) +} + +/// The data fed to one export and the result expected back. +struct HostCase { + /// One tile of values per ABI parameter, in description order; result + /// slots start zeroed. + buffers: Vec>, + /// Index of the parameter whose buffer is read back and checked. + checked: usize, + expected: Vec, + what: &'static str, +} + +/// What the host feeds each export and what it expects back. Every case is +/// exact in f32: relu copies or zeroes, 0.5 halves without rounding, and the +/// warp sums are small integers. +fn host_model(desc: &TileExportDesc, n: usize) -> Fallible { + // A ramp from -4 to 4 crosses zero without landing on it. + let last = (n.max(2) - 1) as f32; + let ramp: Vec = (0..n).map(|i| -4.0 + 8.0 * (i as f32) / last).collect(); + let case = match desc.decl.as_str() { + "relu" => HostCase { + buffers: vec![ramp.clone(), vec![0.0; n]], + checked: 1, + expected: ramp + .iter() + .map(|&x| if x > 0.0 { x } else { 0.0 }) + .collect(), + what: "relu of a ramp from -4 to 4", + }, + "scale_in_place" => HostCase { + buffers: vec![ramp.clone(), vec![0.5; n]], + checked: 0, + expected: ramp.iter().map(|x| x * 0.5).collect(), + what: "the ramp scaled in place by 0.5", + }, + "warp_sum" => { + let lanes: Vec = (0..n).map(|i| i as f32).collect(); + let expected = (0..n) + .map(|i| { + let first = i - i % 32; + lanes[first..(first + 32).min(n)].iter().sum() + }) + .collect(); + HostCase { + buffers: vec![lanes, vec![0.0; n]], + checked: 1, + expected, + what: "32-lane sums over 0..128, every lane holding its warp's total", + } + } + other => return Err(format!("no host model for export `{other}`").into()), + }; + if case.buffers.len() != desc.params.len() || !desc.params[case.checked].access.is_output() { + return Err(format!( + "`{}`: the host model does not match the described parameters", + desc.decl + ) + .into()); + } + Ok(case) +} + +/// Run `call_` over one tile, one thread per footprint, and compare +/// every element of the checked buffer with the host model. +fn check_export( + ctx: &Arc, + module: &Arc, + desc: &TileExportDesc, +) -> Fallible<()> { + let n = usize::try_from(desc.shape[0])?; + let footprint = usize::try_from(desc.params[0].footprint)?; + if n % footprint != 0 { + return Err(format!( + "`{}`: footprint {footprint} does not divide the tile size {n}", + desc.decl + ) + .into()); + } + let invocations = n / footprint; + let case = host_model(desc, n)?; + + let stream = ctx.default_stream(); + let buffers = case + .buffers + .iter() + .map(|data| DeviceBuffer::from_host(&stream, data)) + .collect::, _>>()?; + let mut pointers: Vec<_> = buffers.iter().map(DeviceBuffer::cu_deviceptr).collect(); + let mut args: Vec<*mut c_void> = pointers + .iter_mut() + .map(|pointer| std::ptr::from_mut(pointer).cast::()) + .collect(); + let function = module.load_function(&format!("call_{}", desc.decl))?; + // SAFETY: the generated kernel takes one device pointer per ABI parameter + // in description order, which is how `args` was built; every buffer holds + // `n` elements, one footprint for each of the `n / footprint` threads + // launched, and outlives the launch; the stream belongs to the module's + // context. + unsafe { + launch_kernel_on_stream( + &function, + (1, 1, 1), + (u32::try_from(invocations)?, 1, 1), + 0, + &stream, + &mut args, + )?; + } + + let got = buffers[case.checked].to_host_vec(&stream)?; + let mismatches: Vec<_> = got + .iter() + .zip(&case.expected) + .enumerate() + .filter(|(_, (got, expected))| got != expected) + .collect(); + if let Some((index, (got, expected))) = mismatches.first() { + return Err(format!( + "`{}`: {} of {n} elements differ, first at [{index}]: got {got}, expected {expected}", + desc.decl, + mismatches.len() + ) + .into()); + } + let footprints = if footprint == 1 { + String::new() + } else { + format!(" in {invocations} footprints of {footprint}") + }; + println!( + "PASS {}: {n} elements{footprints} through `{}`, {}", + desc.decl, desc.symbol, case.what + ); + Ok(()) +} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/tile_caller b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/tile_caller new file mode 100755 index 0000000000..b17e0e6b9e Binary files /dev/null and b/cuda-oxide/crates/rustc-codegen-cuda/examples/tilelib_relu/tile_caller differ 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 d4c75a0b72..0162d17bb3 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tma_copy/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/Cargo.toml index 272ed58d2f..575f56a82e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/Cargo.toml @@ -15,4 +15,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/README.md index b89b9503b0..0f8b32dba3 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/README.md @@ -10,9 +10,17 @@ For TMA multicast (also sm_90+), see the [`tma_multicast`](../tma_multicast/) ex ## What This Example Does 1. **tma_copy_2d_test**: Async 2D tile copy from global to shared memory -2. **tma_pipeline_test**: TMA with mbarrier for completion tracking +2. **tma_copy_2d_cta_test**: The same copy with `cp_async_bulk_tensor_2d_g2s_cta`, + whose destination is a `SharedPtr`, the issuing CTA's own shared memory by type +3. **tma_pipeline_test**: TMA with mbarrier for completion tracking -Both kernels use `#[grid_constant] tensor_map: &TmaDescriptor`. The generated +`cp_async_bulk_tensor_2d_g2s` addresses cluster shared memory, so its +destination may belong to another CTA of the cluster. The `_cta` form takes a +`SharedPtr` destination, here `SharedArray::shared_ptr(&raw mut TILE)`, and +emits `cp.async.bulk.tensor.2d.shared::cta.global` without converting any +address; `verify-code-shape.sh` checks both forms in the generated PTX. + +All kernels use `#[grid_constant] tensor_map: &TmaDescriptor`. The generated host launcher takes the descriptor by value and places its 128 bytes directly in the kernel's parameter storage. Every thread can address that same read-only value; no separate device descriptor allocation or upload is needed. The tensor 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 04ba8415de..aa21c06446 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 @@ -9,6 +9,9 @@ //! //! Demonstrates TMA (Tensor Memory Accelerator) usage: //! - cp_async_bulk_tensor_2d_g2s: Async 2D tensor copy global → shared +//! - cp_async_bulk_tensor_2d_g2s_cta: the same copy into the issuing CTA's +//! own shared memory, whose destination is a `SharedPtr`, a CTA-shared +//! pointer by type //! - mbarrier: Barrier-based completion tracking //! - `#[grid_constant]`: the descriptor is passed by value in the launch packet, //! without allocating or uploading a separate descriptor buffer. @@ -34,7 +37,9 @@ use cuda_device::barrier::{ Barrier, fence_proxy_async_shared_cta, mbarrier_arrive, mbarrier_arrive_expect_tx, mbarrier_init, mbarrier_try_wait, }; -use cuda_device::tma::{TmaDescriptor, cp_async_bulk_tensor_2d_g2s}; +use cuda_device::tma::{ + TmaDescriptor, cp_async_bulk_tensor_2d_g2s, cp_async_bulk_tensor_2d_g2s_cta, +}; use cuda_device::{DisjointSlice, SharedArray, kernel, thread}; use cuda_host::cuda_module; use std::mem::MaybeUninit; @@ -125,6 +130,65 @@ mod kernels { } } + /// `tma_copy_2d_test` with the CTA-local copy: the tile and the barrier + /// are this CTA's own shared statics, so the `.shared::cta` form applies. + #[kernel] + pub fn tma_copy_2d_cta_test( + #[grid_constant] tensor_map: &TmaDescriptor, + mut out: DisjointSlice, + tile_x: i32, + tile_y: i32, + ) { + const TILE_SIZE: usize = 64 * 64; + const TILE_BYTES: u32 = (TILE_SIZE * 4) as u32; + static mut TILE: SharedArray = SharedArray::UNINIT; + static mut BAR: Barrier = Barrier::UNINIT; + + let tid = thread::threadIdx_x(); + let block_size = thread::blockDim_x(); + let gid = thread::index_1d(); + + if tid == 0 { + unsafe { + mbarrier_init(&raw mut BAR, block_size); + fence_proxy_async_shared_cta(); + } + } + thread::sync_threads(); + + if tid == 0 { + // SAFETY: TILE and BAR outlive the copy, which completes on BAR + // before any thread reads TILE. + unsafe { + cp_async_bulk_tensor_2d_g2s_cta( + SharedArray::shared_ptr(&raw mut TILE).cast::(), + tensor_map, + tile_x, + tile_y, + &raw mut BAR, + ); + } + } + + let token = unsafe { + if tid == 0 { + mbarrier_arrive_expect_tx(&raw const BAR, 1, TILE_BYTES) + } else { + mbarrier_arrive(&raw const BAR) + } + }; + unsafe { while !mbarrier_try_wait(&raw const BAR, token) {} } + thread::sync_threads(); + + let idx = gid.get(); + if idx < TILE_SIZE { + let val = unsafe { TILE[idx] }; + if let Some(out_elem) = out.get_mut(gid) { + *out_elem = val; + } + } + } + /// Simple TMA pipeline test - ALL threads participate in barrier. #[kernel] pub fn tma_pipeline_test( @@ -227,7 +291,8 @@ fn main() -> Result<(), Box> { println!("✓ Module loaded successfully\n"); // Run tests - run_tma_copy_test(&stream, &module)?; + run_tma_copy_test(&stream, &module, false)?; + run_tma_copy_test(&stream, &module, true)?; run_tma_pipeline_test(&stream, &module)?; println!("\n=== TMA Copy Test Complete ==="); @@ -265,8 +330,14 @@ fn verify_ptx_only(ctx: &Arc) -> Result<(), Box, module: &kernels::LoadedModule, + cta: bool, ) -> Result<(), Box> { - println!("--- Test 1: TMA Copy (tma_copy_2d_test) ---\n"); + let kernel = if cta { + "tma_copy_2d_cta_test" + } else { + "tma_copy_2d_test" + }; + println!("--- Test: TMA Copy ({kernel}) ---\n"); println!( "1. Allocating {}x{} tensor ({} floats, {} KB)", @@ -294,7 +365,7 @@ fn run_tma_copy_test( TILE_HEIGHT, )?; - println!("3. Launching tma_copy_2d_test kernel..."); + println!("3. Launching {kernel} kernel..."); let tile_x: i32 = 0; let tile_y: i32 = 0; @@ -309,14 +380,25 @@ fn run_tma_copy_test( // SAFETY: launch shape/resources match the kernel; buffers cover its accesses. unsafe { - module.tma_copy_2d_test( - (stream).as_ref(), - cfg, - tensor_map, - &mut dev_output, - tile_x, - tile_y, - ) + if cta { + module.tma_copy_2d_cta_test( + (stream).as_ref(), + cfg, + tensor_map, + &mut dev_output, + tile_x, + tile_y, + ) + } else { + module.tma_copy_2d_test( + (stream).as_ref(), + cfg, + tensor_map, + &mut dev_output, + tile_x, + tile_y, + ) + } }?; stream.synchronize()?; @@ -359,7 +441,7 @@ fn run_tma_pipeline_test( stream: &Arc, module: &kernels::LoadedModule, ) -> Result<(), Box> { - println!("\n--- Test 2: TMA Pipeline (tma_pipeline_test) ---\n"); + println!("\n--- Test: TMA Pipeline (tma_pipeline_test) ---\n"); const PIPELINE_TILE_WIDTH: u32 = 32; const PIPELINE_TILE_HEIGHT: u32 = 32; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/verify-code-shape.sh new file mode 100755 index 0000000000..818506f8f9 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_copy/verify-code-shape.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# cp_async_bulk_tensor_2d_g2s_cta must copy into the issuing CTA's shared +# memory: its kernel emits the `.shared::cta` destination form and, since its +# `SharedPtr` destination is shared by type, converts no address to the shared +# or cluster window. The cluster API must keep its +# `.shared::cluster` form, since its destination may belong to another CTA. + +set -euo pipefail + +root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ll="${root}/tma_copy.ll" +ptx="${root}/tma_copy.ptx" + +test -s "${ll}" +test -s "${ptx}" + +# Print the body of one PTX entry. Waits for the header's `{` so a forward +# declaration (which ends in `;`) is skipped. +entry_body() { + local symbol="$1" + awk -v marker="${symbol}(" ' + !emit && !candidate && index($0, marker) && index($0, ".entry") { + candidate = 1 + } + candidate && $0 ~ /^[[:space:]]*;[[:space:]]*$/ { + candidate = 0 + next + } + candidate && $0 ~ /^[[:space:]]*\{[[:space:]]*$/ { + emit = 1 + candidate = 0 + } + emit { print } + emit && index($0, "End function") != 0 { exit } + ' "${ptx}" +} + +require_entry_shape() { + local symbol="$1" + local description="$2" + local pattern="$3" + if ! entry_body "${symbol}" | grep -E "${pattern}" >/dev/null; then + echo "error: missing ${description} in ${ptx}:${symbol}" >&2 + exit 1 + fi +} + +forbid_entry_shape() { + local symbol="$1" + local description="$2" + local pattern="$3" + if entry_body "${symbol}" | grep -E "${pattern}" >/dev/null; then + echo "error: unexpected ${description} in ${ptx}:${symbol}" >&2 + exit 1 + fi +} + +if ! grep -E 'call void @llvm\.nvvm\.cp\.async\.bulk\.tensor\.g2s\.cta\.tile\.2d\(ptr addrspace\(3\)' \ + "${ll}" >/dev/null; then + echo "error: missing typed CTA-local G2S call in ${ll}" >&2 + exit 1 +fi + +require_entry_shape tma_copy_2d_cta_test "CTA-local TMA load" \ + 'cp\.async\.bulk\.tensor\.2d\.shared::cta\.global\.tile\.mbarrier::complete_tx::bytes' +forbid_entry_shape tma_copy_2d_cta_test "cluster-window destination" 'shared::cluster' +forbid_entry_shape tma_copy_2d_cta_test "generic-to-shared address conversion" 'cvta\.to\.shared' + +require_entry_shape tma_copy_2d_test "cluster TMA load" \ + 'cp\.async\.bulk\.tensor\.2d\.shared::cluster\.global\.tile\.mbarrier::complete_tx::bytes' + +echo "tma_copy code shape: PASS" 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 7af886b9bc..41c5da3d69 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tma_multicast/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/Cargo.toml index c37ebdef2e..4bbdbb4fdb 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tma_multicast/Cargo.toml @@ -17,4 +17,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 d07a8d2fa7..fa8576a5c6 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/toolchain_pairing_smoke/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/toolchain_pairing_smoke/Cargo.toml index 82e82efe52..fa14dd4ecf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/toolchain_pairing_smoke/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/toolchain_pairing_smoke/Cargo.toml @@ -14,4 +14,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 3c965c977f..b153af1014 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/Cargo.toml index a4dd1db5c8..3066946c3e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/transmute_scalar_struct/Cargo.toml @@ -19,4 +19,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 ba512ea61c..ad1c253f8c 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tuple_array_provenance/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/Cargo.toml index ea325f4c77..e32bd9e6d5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_array_provenance/Cargo.toml @@ -11,6 +11,6 @@ license = "Apache-2.0" # cargo oxide run tuple_array_provenance [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.lock index a32429c09f..161953c941 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.toml index bb850a76db..3687e797dd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_constant_provenance/Cargo.toml @@ -7,6 +7,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.lock index 44984c38b3..7bd6be0875 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.toml index 58ac6fc674..88a009c6c1 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_field_lookup/Cargo.toml @@ -19,4 +19,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 ec8eee26e6..fa83d0f633 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/tuple_return/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/Cargo.toml index 16b814bef8..d4851155e5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/tuple_return/Cargo.toml @@ -22,4 +22,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 f063fd2267..1ffae9df91 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/unaligned_memory/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/Cargo.toml index adbcd2498b..dc6dcd204c 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unaligned_memory/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.lock index b3ee031b9c..e7e441bcc4 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.toml index 44e8ddb3d8..8f149e7dbd 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing/Cargo.toml @@ -21,5 +21,5 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } ptx-parse = { path = "../../../ptx-parse" } 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 0dae26617c..be29bf5921 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/uninit_const/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/Cargo.toml index 88040199c9..bb1b4de2e9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/uninit_const/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 0ff48d9e4c..12fc11b53d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/union_aggregate/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/Cargo.toml index 40603344dc..b75e4492c5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/union_aggregate/Cargo.toml @@ -14,4 +14,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 5262ec59b0..9583827edc 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/unroll_bounds_check/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/Cargo.toml index a5f46a89e4..4c7f3e6ed7 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_bounds_check/Cargo.toml @@ -9,4 +9,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 8d0cce9016..a234aff07b 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/unroll_smoke/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/Cargo.toml index b2759464d5..6bc89e3a59 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/Cargo.toml @@ -12,4 +12,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 974210636d..2469ed630a 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 @@ -13,7 +13,17 @@ //! unroll it completely and the per-iteration `i & 3` should fold to literals. //! `partial_unroll` has a runtime trip count, so `#[unroll(4)]` unrolls the body //! by 4 and leaves a remainder loop. `range_for_loops` puts the same requests on -//! range `for` loops. All are semantics-preserving; the host checks the sums. +//! range `for` loops, and `signed_range_for_loops` and `signed_partial_range_for` +//! on loops over `i32` counters. `nested_lane_range` is a 16×16 nest over a +//! 16-lane accumulator. A runtime-bounded mix in the inner range is what keeps +//! LLVM from unrolling that nest itself. +//! `isize_range_for` covers `a..b`, `(a..b).rev()`, and `#[unroll(4)]` on an +//! `isize` counter. +//! All are semantics-preserving; the host checks the sums. +//! +//! `verify-code-shape.sh` asserts that `nested_lane_range` and `isize_range_for` +//! have no local memory in the generated PTX. GPU-less CI runs it through +//! `scripts/smoketest.sh --compile-only`. //! //! Run: cargo oxide run unroll_smoke @@ -504,6 +514,129 @@ mod kernels { *partial = sum; *unrecognized = other; } + + /// Range `for` loops over signed counters. A counter used only through `as` + /// has no type to infer from, so it is `i32`, and core steps a signed counter + /// through `checked_add` (`checked_sub` for `.rev()`) rather than a plain + /// add. `#[unroll]` on `0..16` makes the `acc` indices literals, so the array + /// stays in registers: the PTX holds no local-memory access and no loop + /// branch for it. `acc[j]` sums `input[16 * r + j]` over four rounds, and + /// `#[unroll]` on `(0..16).rev()` folds `acc` into `sum(acc[j] * 3^j)`. + #[kernel] + pub fn signed_range_for_loops(input: &[u32], mut folded: DisjointSlice) { + let base = thread::index_1d().get() as u32; + let Some(folded) = folded.get_mut(thread::index_1d()) else { + return; + }; + let mut acc = [0u32; 16]; + let mut round = 0usize; + while round < 4 { + #[unroll] + for j in 0..16 { + // SAFETY: the host passes 64 elements. + let x = unsafe { *input.get_unchecked(round * 16 + j as usize) }; + acc[j as usize] = acc[j as usize].wrapping_add(x); + } + round += 1; + } + let mut sum = 0u32; + #[unroll] + for j in (0..16).rev() { + sum = sum.wrapping_mul(3).wrapping_add(acc[j as usize]); + } + *folded = sum.wrapping_add(base); + } + + /// `#[unroll(4)]` on `0..n` with `n: i32` sums `n*(n-1)/2`. + #[kernel] + pub fn signed_partial_range_for(mut partial: DisjointSlice, n: i32) { + let Some(partial) = partial.get_mut(thread::index_1d()) else { + return; + }; + let mut sum = 0u32; + #[unroll(4)] + for i in 0..n { + sum = sum.wrapping_add(i as u32); + } + *partial = sum; + } + + /// 16×16 nested range over a 16-lane accumulator. Each loaded value is mixed + /// `steps` times. That runtime-bounded loop is enough for LLVM to leave the + /// nest looped, with `acc` in a `.local` depot (`ld.local` / `st.local`). + /// `#[unroll]` on both range loops makes the lane indices literals, so `acc` + /// stays in registers. `acc[j]` sums the mixed `input[i * 16 + j]` over `i` + /// in `0..16`, and the result is that vector summed, plus the thread index. + #[allow(clippy::needless_range_loop)] + #[kernel] + pub fn nested_lane_range(input: &[u32], mut out: DisjointSlice, steps: u32) { + let base = thread::index_1d().get() as u32; + let Some(out_elem) = out.get_mut(thread::index_1d()) else { + return; + }; + let mut acc = [0u32; 16]; + #[unroll] + for i in 0..16 { + #[unroll] + for j in 0..16 { + // SAFETY: the host passes 256 elements, and both loops cover `0..16`. + let mut x = unsafe { *input.get_unchecked(i * 16 + j) }; + let mut k = 0u32; + while k < steps { + x = x.wrapping_mul(3).wrapping_add(1); + k += 1; + } + acc[j] = acc[j].wrapping_add(x); + } + } + let mut sum = 0u32; + #[unroll] + for lane in 0..16 { + sum = sum.wrapping_add(acc[lane]); + } + *out_elem = sum.wrapping_add(base); + } + + /// Range `for` loops over `isize`. core steps `isize` by casting to `i64` + /// and back; those are the same dialect type, so the loops unroll like + /// `i64`. `#[unroll]` on `0isize..8` (`a..b`) fills `acc`, `#[unroll]` on + /// `(0isize..8).rev()` folds it, and `#[unroll(4)]` on `0isize..n` sums + /// `n*(n-1)/2`. + #[allow(clippy::needless_range_loop)] + #[kernel] + pub fn isize_range_for( + input: &[u32], + mut folded: DisjointSlice, + mut partial: DisjointSlice, + n: isize, + ) { + let base = thread::index_1d().get() as u32; + let (Some(folded), Some(partial)) = ( + folded.get_mut(thread::index_1d()), + partial.get_mut(thread::index_1d()), + ) else { + return; + }; + let mut acc = [0u32; 8]; + #[unroll] + for i in 0isize..8 { + // SAFETY: the host passes at least 8 elements. + let x = unsafe { *input.get_unchecked(i as usize) }; + acc[i as usize] = x; + } + let mut weighted = 0u32; + #[unroll] + for i in (0isize..8).rev() { + weighted = weighted.wrapping_mul(3).wrapping_add(acc[i as usize]); + } + let mut sum = 0u32; + #[unroll(4)] + for i in 0isize..n { + sum = sum.wrapping_add(i as u32); + } + *folded = weighted.wrapping_add(base); + *partial = sum; + } } fn main() { @@ -711,6 +844,111 @@ fn main() { vec![28; N], "`0..=4` and `(0..10).step_by(3)`" ); + + let signed_input: Vec = (0..64u32).map(|i| i * 7 + 1).collect(); + let d_signed_input = DeviceBuffer::from_host(&stream, &signed_input).unwrap(); + let mut d_signed = DeviceBuffer::::zeroed(&stream, N).unwrap(); + let mut d_signed_part = DeviceBuffer::::zeroed(&stream, N).unwrap(); + // SAFETY: each thread writes its own element of the buffer; the kernel + // reads the 64 input elements. + unsafe { + module.signed_range_for_loops(stream.as_ref(), cfg, &d_signed_input, &mut d_signed) + } + .expect("launch signed_range_for_loops"); + // SAFETY: each thread writes its own element of the buffer. + unsafe { + module.signed_partial_range_for( + stream.as_ref(), + cfg, + &mut d_signed_part, + range_trip as i32, + ) + } + .expect("launch signed_partial_range_for"); + let want_folded: Vec = (0..N as u32) + .map(|tid| { + (0..16) + .map(|j| (0..4).fold(0u32, |sum, r| sum.wrapping_add(signed_input[16 * r + j]))) + .rev() + .fold(0u32, |sum, acc| sum.wrapping_mul(3).wrapping_add(acc)) + .wrapping_add(tid) + }) + .collect(); + assert_eq!( + d_signed.to_host_vec(&stream).unwrap(), + want_folded, + "`i32` `0..16` and `(0..16).rev()` over a local array" + ); + assert_eq!( + d_signed_part.to_host_vec(&stream).unwrap(), + vec![range_trip * (range_trip - 1) / 2; N], + "partial `i32` `0..{range_trip}`" + ); + + let nested_input: Vec = (0..256u32) + .map(|i| i.wrapping_mul(3).wrapping_add(1)) + .collect(); + let d_nested_input = DeviceBuffer::from_host(&stream, &nested_input).unwrap(); + let mut d_nested = DeviceBuffer::::zeroed(&stream, N).unwrap(); + // SAFETY: each thread writes its own element; the kernel reads 256 inputs. + let steps = 1u32; + unsafe { + module.nested_lane_range(stream.as_ref(), cfg, &d_nested_input, &mut d_nested, steps) + } + .expect("launch nested_lane_range"); + let lane_sum = (0..16).fold(0u32, |sum, lane| { + sum.wrapping_add((0..16).fold(0u32, |acc, row| { + let mut x = nested_input[row * 16 + lane]; + for _ in 0..steps { + x = x.wrapping_mul(3).wrapping_add(1); + } + acc.wrapping_add(x) + })) + }); + let want_nested: Vec = (0..N as u32) + .map(|tid| lane_sum.wrapping_add(tid)) + .collect(); + assert_eq!( + d_nested.to_host_vec(&stream).unwrap(), + want_nested, + "16x16 nested range over a 16-lane accumulator" + ); + + let mut d_isize = DeviceBuffer::::zeroed(&stream, N).unwrap(); + let mut d_isize_part = DeviceBuffer::::zeroed(&stream, N).unwrap(); + // SAFETY: each thread writes its own element of each buffer; the kernel + // reads the first 8 input elements. + unsafe { + module.isize_range_for( + stream.as_ref(), + cfg, + &d_nested_input, + &mut d_isize, + &mut d_isize_part, + range_trip as isize, + ) + } + .expect("launch isize_range_for"); + let want_isize: Vec = (0..N as u32) + .map(|tid| { + (0..8) + .rev() + .fold(0u32, |sum, i| { + sum.wrapping_mul(3).wrapping_add(nested_input[i]) + }) + .wrapping_add(tid) + }) + .collect(); + assert_eq!( + d_isize.to_host_vec(&stream).unwrap(), + want_isize, + "`isize` `0..8` and `(0..8).rev()`" + ); + assert_eq!( + d_isize_part.to_host_vec(&stream).unwrap(), + vec![range_trip * (range_trip - 1) / 2; N], + "partial `isize` `0..{range_trip}`" + ); } let mut failures = 0usize; diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/verify-code-shape.sh new file mode 100755 index 0000000000..035c8e0c5e --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/unroll_smoke/verify-code-shape.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +set -euo pipefail + +root="$(cd "$(dirname "$0")" && pwd)" +ptx="${1:-$root/unroll_smoke.ptx}" + +if [[ ! -f "$ptx" ]]; then + echo "PTX file not found: $ptx" >&2 + exit 1 +fi + +python3 - "$ptx" <<'PY' +import re +import sys + +ptx_path = sys.argv[1] +ptx = open(ptx_path, encoding="utf-8").read() +# Strip comments so a mention in a comment cannot satisfy or hide a check. +ptx = re.sub(r"/\*.*?\*/|//[^\n]*", "", ptx, flags=re.S) +# Kernels whose accumulator must stay in registers. A `.local` depot here +# means the annotated range loops were left for LLVM, which does not unroll +# `nested_lane_range`. +names = ("nested_lane_range", "isize_range_for") +local_memory = re.compile(r"(?m)(?:^|\s)\.local\b|ld\.local|st\.local") +for name in names: + header = re.search(r"\.entry\s+" + re.escape(name) + r"\s*\(", ptx) + if header is None: + sys.exit(f"error: missing PTX entry {name} in {ptx_path}") + brace = ptx.find("{", header.end()) + if brace < 0: + sys.exit(f"error: entry {name} has no body") + depth = 1 + end = brace + 1 + while depth and end < len(ptx): + depth += (ptx[end] == "{") - (ptx[end] == "}") + end += 1 + if depth: + sys.exit(f"error: unterminated PTX entry {name}") + body = ptx[brace + 1 : end - 1] + if local_memory.search(body): + sys.exit(f"error: {name} uses local memory in {ptx_path}") + print(f"{name}: no local memory PASS") +print("unroll_smoke PTX shape: PASS") +PY 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 dead9d09ff..2e6d5cb154 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/vecadd/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.toml index 139c876e8d..d74ce061f5 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vecadd/Cargo.toml @@ -22,4 +22,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 07061c7cb4..ecc9160ea6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -717,3 +747,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/vectorization/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.toml index 6e192f3ef3..0781724d84 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/vectorization/Cargo.toml @@ -21,5 +21,5 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } ptx-parse = { path = "../../../ptx-parse" } 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 3fa9c5e123..255b16ace7 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 @@ -155,6 +155,136 @@ mod view_kernels { empty.as_mut_ptr() } + /// Regression coverage for a 16-byte enum filler beside a pointer niche. + /// + /// The `[u64; 2]` payload occupies 16 bytes before the slice pointer. Enum + /// storage must represent that filler without obscuring the slice pointer's + /// global address space; otherwise the final vector load becomes generic + /// `ld.v4.*` instead of `ld.global.v4.*`. + #[kernel] + pub fn f32x4_view_copy_through_wide_option( + input: &[f32], + mut output: DisjointSlice, + lo: u64, + hi: u64, + ) { + let i = thread::index_1d().get(); + + let out_len = output.len(); + let out_flat = unsafe { core::slice::from_raw_parts_mut(output.as_mut_ptr(), out_len) }; + + let Some(input_quads) = vector::as_vectors::(input) else { + return; + }; + + let Some(out_quads) = vector::as_vectors_mut::(out_flat) else { + return; + }; + + let cookie = [lo, hi]; + + let wrapped = if lo == u64::MAX { + None + } else { + Some((cookie, input_quads)) + }; + + let Some((_cookie, quads)) = wrapped else { + return; + }; + + if i < quads.len() && i < out_quads.len() { + out_quads[i] = quads[i]; + } + } + + /// Regression coverage for a 16-byte enum filler beside a mutable pointer niche. + /// + /// This mirrors `f32x4_view_copy_through_wide_option`, but places the output + /// slice behind the wide enum payload. Losing its address space would turn the + /// vector store into generic `st.v4.*` instead of `st.global.v4.*`. + #[kernel] + pub fn f32x4_view_store_through_wide_option( + input: &[f32], + mut output: DisjointSlice, + lo: u64, + hi: u64, + ) { + let i = thread::index_1d().get(); + + let Some(input_quads) = vector::as_vectors::(input) else { + return; + }; + + let out_len = output.len(); + let out_flat = unsafe { core::slice::from_raw_parts_mut(output.as_mut_ptr(), out_len) }; + + let Some(out_quads) = vector::as_vectors_mut::(out_flat) else { + return; + }; + + let cookie = [lo, hi]; + + let wrapped = if lo == u64::MAX { + None + } else { + Some((cookie, out_quads)) + }; + + let Some((_cookie, out_quads)) = wrapped else { + return; + }; + + if i < input_quads.len() && i < out_quads.len() { + out_quads[i] = input_quads[i]; + } + } + + /// Regression coverage for a 24-byte enum filler beside a pointer niche. + /// + /// `[u64; 3]` leaves 24 unclaimed bytes before the slice pointer. Typed + /// as `[3 x i64]` the enum storage is 5 leaves, under the block-argument + /// scalarizer's budget of 16, so the `Some`/`None` join splits into + /// scalar phis. As `[24 x i8]` it was 26 leaves, the join stayed one + /// aggregate phi, and the vector load lost its global window the same + /// way the 16-byte case did. + #[kernel] + pub fn f32x4_view_copy_through_24_byte_option( + input: &[f32], + mut output: DisjointSlice, + lo: u64, + hi: u64, + ) { + let i = thread::index_1d().get(); + + let out_len = output.len(); + let out_flat = unsafe { core::slice::from_raw_parts_mut(output.as_mut_ptr(), out_len) }; + + let Some(input_quads) = vector::as_vectors::(input) else { + return; + }; + + let Some(out_quads) = vector::as_vectors_mut::(out_flat) else { + return; + }; + + let cookie = [lo, hi, lo ^ hi]; + + let wrapped = if lo == u64::MAX { + None + } else { + Some((cookie, input_quads)) + }; + + let Some((_cookie, quads)) = wrapped else { + return; + }; + + if i < quads.len() && i < out_quads.len() { + out_quads[i] = quads[i]; + } + } + /// Copy one `F32x4` quad per thread between flat `f32` buffers through /// checked [`vector::as_vectors`] views. The quad is moved whole, never /// decomposed into lanes, so both the load and the store fuse into @@ -250,6 +380,53 @@ fn main() { .expect("f32x4_view_copy launch"); let flat_round_trip = flat_out.to_host_vec(&stream).expect("flat readback") == flat; + // The wide-option kernels route the same copy through + // `Option<([u64; 2], &[F32x4])>` and `Option<([u64; 3], &[F32x4])>`: a + // 16- or 24-byte payload beside the slice pointer. Launch each into a + // fresh zeroed buffer and check the round trip, then take the `None` arm + // once (`lo == u64::MAX`) and check that the output stays untouched. + let wide_option_output = |launch: &dyn Fn(&mut DeviceBuffer)| -> Vec { + let mut out = DeviceBuffer::::zeroed(&stream, N * 4).expect("wide-option output"); + launch(&mut out); + out.to_host_vec(&stream).expect("wide-option readback") + }; + // SAFETY (each launch below): launch shape/resources match the kernel; + // buffers cover its accesses. + let wide_copy_round_trip = wide_option_output(&|out| { + unsafe { + view_module.f32x4_view_copy_through_wide_option(&stream, cfg, &flat_in, out, 123, 456) + } + .expect("f32x4_view_copy_through_wide_option launch"); + }) == flat; + let wide_store_round_trip = wide_option_output(&|out| { + unsafe { + view_module.f32x4_view_store_through_wide_option(&stream, cfg, &flat_in, out, 123, 456) + } + .expect("f32x4_view_store_through_wide_option launch"); + }) == flat; + let wide_none_untouched = wide_option_output(&|out| { + unsafe { + view_module.f32x4_view_copy_through_wide_option( + &stream, + cfg, + &flat_in, + out, + u64::MAX, + 456, + ) + } + .expect("f32x4_view_copy_through_wide_option launch, None arm"); + }) + .iter() + .all(|v| *v == 0.0); + let option_24_round_trip = wide_option_output(&|out| { + unsafe { + view_module + .f32x4_view_copy_through_24_byte_option(&stream, cfg, &flat_in, out, 123, 456) + } + .expect("f32x4_view_copy_through_24_byte_option launch"); + }) == flat; + // Inspect the code we just launched: read the PTX back from the artifact // bundle embedded in this binary. Those are the same bytes `kernels::load` // handed to the driver, so the shape assertion can never diverge from what @@ -269,6 +446,22 @@ fn main() { errors += 1; println!(" !! f32x4_view_copy round-trip copy mismatch"); } + if !wide_copy_round_trip { + errors += 1; + println!(" !! f32x4_view_copy_through_wide_option round-trip copy mismatch"); + } + if !wide_store_round_trip { + errors += 1; + println!(" !! f32x4_view_store_through_wide_option round-trip copy mismatch"); + } + if !wide_none_untouched { + errors += 1; + println!(" !! f32x4_view_copy_through_wide_option wrote output through the None arm"); + } + if !option_24_round_trip { + errors += 1; + println!(" !! f32x4_view_copy_through_24_byte_option round-trip copy mismatch"); + } if let Some(ptx) = &ptx { let document = ptx_parse::Document::parse(ptx).expect("parse embedded PTX"); @@ -326,6 +519,32 @@ fn main() { } } } + + // A wide enum payload next to the slice pointer must not hide the + // pointer's global address space from NVPTX address-space inference. + // The enum storage has to stay under the block-argument scalarizer's + // leaf budget: a typed filler (`[2 x i64]`, `[3 x i64]`) does, a byte + // array (`[16 x i8]`, `[24 x i8]`) does not, and then the merged + // pointer rides an aggregate phi that `InferAddressSpaces` cannot + // see through. The copy kernels take the input slice through the + // enum (load side); the store kernel takes the output slice. + for (kernel, dir, filler) in [ + ("f32x4_view_copy_through_wide_option", "ld", "16-byte"), + ("f32x4_view_store_through_wide_option", "st", "16-byte"), + ("f32x4_view_copy_through_24_byte_option", "ld", "24-byte"), + ] { + let body = kernel_body(&document, kernel); + match find_128bit_global(body, dir) { + Some(line) => println!("{kernel}: {line}"), + None => { + errors += 1; + let access = if dir == "ld" { "load" } else { "store" }; + println!( + " !! {kernel}: {filler} enum filler lost the 128-bit global {access}" + ); + } + } + } } else { println!( "PTX shape check skipped: the embedded bundle has no PTX text to \ @@ -386,12 +605,22 @@ fn kernel_body<'source>(document: &ptx_parse::Document<'source>, name: &str) -> /// (`ld.global.v4.f32` on older llc, `ld.global.v2.b64` on llc-21/22); every /// accepted form is a single 128-bit access. fn find_128bit_global<'a>(body: &'a str, dir: &str) -> Option<&'a str> { - let prefix = format!("{dir}.global.v"); + let prefix = format!("{dir}.global."); body.lines().map(str::trim).find(|t| { - t.starts_with(&prefix) - && [".v4.f32", ".v4.b32", ".v2.f64", ".v2.b64"] - .iter() - .any(|w| t.contains(w)) + let Some(rest) = t.strip_prefix(&prefix) else { + return false; + }; + // Read-only inputs may use ld.global.nc; its cache choice preserves + // the global address space and transaction width checked here. + let rest = if dir == "ld" { + rest.strip_prefix("nc.").unwrap_or(rest) + } else { + rest + }; + matches!( + rest.split_whitespace().next(), + Some("v4.f32" | "v4.b32" | "v2.f64" | "v2.b64") + ) }) } 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 31a9923785..8d9e203b9d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/volatile_memory/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/Cargo.toml index b5a9a20356..901bc7d932 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/volatile_memory/Cargo.toml @@ -10,4 +10,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 6225903e18..2dc4af0e91 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/warp_reduce/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/Cargo.toml index 9819a3f4b0..4f9d544909 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_reduce/Cargo.toml @@ -20,4 +20,4 @@ cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } # Host-side CUDA runtime -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 cfffc70730..45f571902d 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/warp_sums/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_sums/Cargo.toml index 6fcb132ad7..7682de36ee 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_sums/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/warp_sums/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock index 288c97ac06..f0236ae67e 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.lock @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/wgmma/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.toml index 1988774e68..7ce56a42cf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma/Cargo.toml @@ -15,4 +15,4 @@ license = "Apache-2.0" [dependencies] cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } 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 2fe046b8aa..e54e6354c7 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 @@ -128,7 +128,7 @@ dependencies = [ [[package]] name = "cuda-bindings" -version = "0.4.0" +version = "0.5.0" dependencies = [ "bindgen", "libloading", @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "cuda-core" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "cuda-bindings", @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "cuda-core-derive" -version = "0.4.0" +version = "0.5.0" dependencies = [ "proc-macro2", "quote", @@ -171,6 +171,7 @@ version = "0.2.1" dependencies = [ "cuda-artifact-finalizer", "cuda-core", + "cuda-device", "cuda-macros", "half", "libc", @@ -184,6 +185,7 @@ dependencies = [ name = "cuda-macros" version = "0.2.1" dependencies = [ + "oxide-tile-export", "proc-macro2", "quote", "reserved-oxide-symbols", @@ -302,6 +304,12 @@ 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" @@ -408,6 +416,15 @@ dependencies = [ "object", ] +[[package]] +name = "oxide-tile-export" +version = "0.2.1" +dependencies = [ + "serde", + "serde_json", + "thiserror", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -522,6 +539,19 @@ dependencies = [ "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" @@ -716,3 +746,9 @@ dependencies = [ "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/crates/rustc-codegen-cuda/examples/wgmma_mma_bf16/Cargo.toml b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma_mma_bf16/Cargo.toml index 880eb226af..13d948ffe4 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma_mma_bf16/Cargo.toml +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/wgmma_mma_bf16/Cargo.toml @@ -8,6 +8,6 @@ license = "Apache-2.0" [workspace] [dependencies] -cuda-core = { path = "../../../../../cuda-core", version = "0.4.0" } +cuda-core = { path = "../../../../../cuda-core" } cuda-device = { path = "../../../cuda-device" } cuda-host = { path = "../../../cuda-host" } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs index c361689734..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 @@ -246,10 +385,11 @@ enum CollectDecision { // match. The mutual-exclusion guarantee between `DEVICE_PREFIX` and // `DEVICE_EXTERN_PREFIX` means we no longer need the historical // "test extern first" ordering dance that lived here previously. +use oxide_tile_export::TileExportDesc; use reserved_oxide_symbols::{ - device_extern_base_name, is_current_device_symbol, is_current_kernel_symbol, + device_base_name, device_extern_base_name, is_current_device_symbol, is_current_kernel_symbol, is_device_extern_symbol, is_device_symbol, is_kernel_symbol, is_legacy_kernel_symbol, - is_ptx_merge_required_marker, kernel_base_name, + is_ptx_merge_required_marker, is_tile_export_meta_symbol, kernel_base_name, }; /// Sanitize a symbol name for use as a PTX identifier. @@ -396,6 +536,12 @@ pub struct CollectionResult<'tcx> { /// dependency. Cubin materialization must reject this case until it can /// reproduce the same cross-crate linking semantics. pub requires_ptx_bundle_merge: bool, + + /// `#[tile_export]` descriptions read from this crate's metadata + /// functions, sorted by implementation symbol. Each one's implementation + /// is a root in `functions`; the backend writes them to + /// `.tileexports.json` beside the device artifact. + pub tile_exports: Vec, } /// Counts kernel functions across all codegen units. @@ -916,6 +1062,145 @@ impl DiscoveryCtx { } } +/// Read every `#[tile_export]` description embedded in this crate. +/// +/// The macro emits one host-only `#[no_mangle] fn +/// cuda_oxide_tile_export__() -> &'static str` per export whose +/// body returns the JSON description. That shape needs no file IO or side +/// channel: `#[no_mangle]` items are always codegened, so the function is in +/// the CGUs even though nothing calls it, and its description is a constant +/// in MIR. The function is never compiled for the device; see the explicit +/// skips in [`collect_device_functions`] and `process_call_operand`. +/// +/// Malformed metadata is fatal rather than skipped: a description that does +/// not parse or validate means the macro and this backend disagree, and a +/// silently missing export would surface much later as a tile library +/// without the symbol the caller asked for. +fn collect_tile_exports<'tcx>( + tcx: TyCtxt<'tcx>, + cgus: &[CodegenUnit<'tcx>], + verbose: bool, +) -> Vec { + let mut descs: Vec = Vec::new(); + let mut seen = HashSet::new(); + for cgu in cgus { + for (item, _data) in cgu.items() { + let MonoItem::Fn(instance) = item else { + continue; + }; + let def_id = instance.def_id(); + let def_path = tcx.def_path_str(def_id); + if !is_tile_export_meta_symbol(&def_path) { + continue; + } + // Partitioning can place one item in several CGUs. + if !seen.insert(def_id) { + continue; + } + let Some(text) = tile_export_meta_text(tcx, *instance) else { + tcx.dcx().span_fatal( + tcx.def_span(def_id), + format!( + "[rustc_codegen_cuda] `{def_path}` carries the reserved tile export \ + metadata prefix but does not return one constant string; only \ + `#[tile_export]` may define functions with this prefix" + ), + ); + }; + let desc = TileExportDesc::from_json(&text).unwrap_or_else(|error| { + tcx.dcx().span_fatal( + tcx.def_span(def_id), + format!( + "[rustc_codegen_cuda] tile export metadata in `{def_path}` is not a \ + description this backend understands ({error}); rebuild with \ + cuda-macros from the same cuda-oxide revision as the backend" + ), + ) + }); + if let Err(error) = desc.validate() { + tcx.dcx().span_fatal( + tcx.def_span(def_id), + format!( + "[rustc_codegen_cuda] tile export `{}` is invalid: {error}", + desc.decl + ), + ); + } + if let Some(previous) = descs.iter().find(|other| other.symbol == desc.symbol) { + tcx.dcx().span_fatal( + tcx.def_span(def_id), + format!( + "[rustc_codegen_cuda] tile exports `{}` and `{}` both use the \ + implementation symbol `{}`; give one of them a distinct `symbol`", + previous.decl, desc.decl, desc.symbol + ), + ); + } + // The declaration name is what a cuTile caller resolves, so two + // exports sharing one would collide in the tile library; catch it + // here rather than when `cargo oxide tilelib` builds the modules. + if let Some(previous) = descs.iter().find(|other| other.decl == desc.decl) { + tcx.dcx().span_fatal( + tcx.def_span(def_id), + format!( + "[rustc_codegen_cuda] tile exports `{}` and `{}` both declare the \ + name `{}` to cuTile; give one of them a distinct `name`", + previous.symbol, desc.symbol, desc.decl + ), + ); + } + if verbose { + eprintln!( + "[collector] Found tile export: {} -> {} (shape {:?}, {} params, sync {})", + desc.decl, + desc.symbol, + desc.shape, + desc.params.len(), + desc.sync.tile_ir_name() + ); + } + descs.push(desc); + } + } + // Stable output regardless of CGU partitioning. + descs.sort_by(|a, b| a.symbol.cmp(&b.symbol)); + descs +} + +/// The constant `&str` a tile export metadata function returns. +/// +/// The optimized body is `_0 = const ""; return`, so the description is +/// the constant operand assigned to `RETURN_PLACE`. The constant evaluates to +/// a `ConstValue::Slice` (or `Indirect`) of UTF-8 bytes; `const_str_text` +/// reads both representations. Anything else (no such assignment, a +/// non-constant, or more than one) returns `None` and the caller reports it. +fn tile_export_meta_text<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Option { + use rustc_middle::mir::{Operand, RETURN_PLACE, Rvalue, StatementKind}; + + let mir = tcx.instance_mir(instance.def); + let mut found = None; + for block in mir.basic_blocks.iter() { + for statement in &block.statements { + let StatementKind::Assign(assign) = &statement.kind else { + continue; + }; + let (place, rvalue) = &**assign; + if place.as_local() != Some(RETURN_PLACE) { + continue; + } + // The second field is the retag marker, irrelevant to a constant. + let Rvalue::Use(Operand::Constant(constant), _) = rvalue else { + return None; + }; + let text = const_str_text(tcx, constant)?; + if found.replace(text).is_some() { + return None; + } + } + } + found +} + /// Collects all device-reachable functions starting from kernel entry points. /// /// This is the main entry point for device function collection. It: @@ -941,6 +1226,7 @@ pub fn collect_device_functions<'tcx>( verbose: bool, ) -> CollectionResult<'tcx> { let mut collector = DeviceCollector::new(tcx, verbose); + let tile_exports = collect_tile_exports(tcx, cgus, verbose); let mut requires_ptx_bundle_merge = cgus.iter().any(|cgu| { cgu.items().iter().any(|(item, _data)| match item { MonoItem::Static(def_id) => is_ptx_merge_required_marker(&tcx.def_path_str(*def_id)), @@ -948,10 +1234,13 @@ pub fn collect_device_functions<'tcx>( }) }); - // Find all kernel entry points + // Find all kernel entry points. `#[tile_export]` metadata functions are + // host-only (their descriptions were read above) and never device code; + // the prefix cannot match a kernel, but say so here rather than rely on it. for cgu in cgus { for (item, _data) in cgu.items() { if let MonoItem::Fn(instance) = item + && !is_tile_export_meta_symbol(&tcx.def_path_str(instance.def_id())) && is_kernel_function(tcx, instance.def_id()) { // Skip generic (non-monomorphized) instances. @@ -1005,15 +1294,67 @@ pub fn collect_device_functions<'tcx>( } } + let has_kernels = !collector.worklist.is_empty(); + + // `#[tile_export]` implementations are roots whether or not the crate has + // kernels: no device code calls them, and the standalone scan below only + // runs for kernel-less crates. The description's `symbol` is the base + // name of the `#[device]` wrapper, which is also what the symbol becomes + // in the device artifact once llvm-export strips the device prefix. + for desc in &tile_exports { + let mut found = false; + for cgu in cgus { + for (item, _data) in cgu.items() { + let MonoItem::Fn(instance) = item else { + continue; + }; + let raw_name = tcx.def_path_str(instance.def_id()); + if device_base_name(&raw_name) != Some(desc.symbol.as_str()) + || !is_device_function(tcx, instance.def_id()) + || !is_fully_monomorphized(tcx, *instance) + // The macro rejects generics, so a generic device function + // sharing the base name is not the implementation. + || !instance.args.is_empty() + { + continue; + } + let name = collector.fqdn(*instance); + let export_name = collector.compute_export_name(&name, *instance); + if verbose { + eprintln!( + "[collector] Found tile export implementation: {} -> {}", + name, export_name + ); + } + collector.add_root(*instance, false, export_name); + found = true; + } + } + if !found { + tcx.dcx().fatal(format!( + "[rustc_codegen_cuda] tile export `{}` names the implementation symbol `{}`, \ + but this crate has no `#[device]` function with that name. `#[tile_export]` \ + generates the implementation next to the exported function; keep the two \ + together and do not rename the generated item", + desc.decl, desc.symbol + )); + } + } + // Find standalone device function roots (Phase 2: device functions without kernels). // Only scan when there are NO kernels — when kernels exist, device functions are - // already collected transitively via the call graph walk. - if collector.worklist.is_empty() { + // already collected transitively via the call graph walk. Tile export + // implementations were rooted above and are skipped here so their export + // names are not computed twice. + if !has_kernels { for cgu in cgus { for (item, _data) in cgu.items() { if let MonoItem::Fn(instance) = item + // Host-only metadata, never device code (see the kernel scan). + && !is_tile_export_meta_symbol(&tcx.def_path_str(instance.def_id())) && is_device_function(tcx, instance.def_id()) && is_fully_monomorphized(tcx, *instance) + && !collector.is_seen(*instance) { let raw_name = tcx.def_path_str(instance.def_id()); @@ -1045,6 +1386,7 @@ pub fn collect_device_functions<'tcx>( // Process the worklist to collect all reachable functions let mut result = collector.collect(); result.requires_ptx_bundle_merge = requires_ptx_bundle_merge; + result.tile_exports = tile_exports; result } @@ -1120,6 +1462,11 @@ impl<'tcx> DeviceCollector<'tcx> { ) } + /// True once `instance` has been rooted or discovered. + fn is_seen(&self, instance: Instance<'tcx>) -> bool { + self.seen.contains(self.tcx.symbol_name(instance).name) + } + /// Adds a root function (kernel) to start collection from. fn add_root(&mut self, instance: Instance<'tcx>, is_kernel: bool, export_name: String) { // Use mangled name as the unique key - this distinguishes different @@ -1177,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); } @@ -1193,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. @@ -1254,6 +1601,7 @@ impl<'tcx> DeviceCollector<'tcx> { functions: self.result, device_externs: self.device_externs, requires_ptx_bundle_merge: false, + tile_exports: Vec::new(), } } @@ -1498,6 +1846,28 @@ impl<'tcx> DeviceCollector<'tcx> { return; }; let fn_path = self.tcx.def_path_str(*def_id); + + // `#[tile_export]` metadata functions are host-only: their body is + // the JSON description the collector already consumed. A device-side + // call means user code reached for the generated item. Refuse it + // here instead of skipping: a skipped call would leave a dangling + // symbol that fails later at module verification. + if is_tile_export_meta_symbol(&fn_path) { + self.tcx + .dcx() + .struct_span_fatal( + call_span, + format!( + "device code in `{}` calls `{fn_path}`, a `#[tile_export]` metadata \ + function; it is host-only and holds the export's JSON description, \ + not device code", + caller.export_name + ), + ) + .with_help("call the exported function itself, or its `` wrapper") + .emit() + } + if fn_path.contains("DynamicSharedArray") && (fn_path.contains("::get") || fn_path.contains("::get_raw") @@ -1684,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; } @@ -1747,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) { @@ -1801,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!( @@ -1912,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 => { @@ -2219,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( @@ -2464,10 +2918,9 @@ impl<'tcx> DeviceCollector<'tcx> { /// "Symbol ... not found" error. /// /// Functions that are declared diverging (return type `!`) are left - /// alone: their call sites have no target block, so the translator - /// already lowers them to LLVM `unreachable` (note: NOT a trap; - /// the optimizer may delete paths that provably reach it, a known - /// gap tracked separately). + /// alone. If their body is emitted, the importer preserves the call and + /// terminates its caller with `unreachable`; if collection deliberately + /// omits the body, the importer retains the existing device-trap fallback. fn check_unreachable_callee( &self, resolved: Instance<'tcx>, @@ -2547,8 +3000,9 @@ impl<'tcx> DeviceCollector<'tcx> { // No stub marker: leave non-local functions to the existing silent // skip (this is what keeps the real `cuda_device` intrinsic - // placeholders and `core`'s cold panic wrappers working), and leave - // declared-diverging functions to the translator's `unreachable` lowering. + // placeholders and `core`'s cold panic wrappers working). Declared- + // diverging functions that are omitted here retain the translator's + // device-trap fallback rather than producing an undefined call. if !def_id.is_local() || mir.return_ty().is_never() { return; } @@ -2577,7 +3031,7 @@ impl<'tcx> DeviceCollector<'tcx> { .with_help( "annotate device helpers with `#[device]`; if the panic is \ intentional, declare the function as diverging (`-> !`) so the \ - call lowers to LLVM `unreachable`", + device pipeline can treat the path as non-returning", ) .emit() } @@ -2748,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, @@ -2871,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 7e62630445..f22d0937d9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/device_codegen.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/device_codegen.rs @@ -349,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. @@ -364,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, @@ -849,6 +855,29 @@ impl From for DeviceCodegenError { } } +/// Whether this `#[inline(never)]` is the one `#[device]` puts on the prefixed +/// function of a device function with its own type or const parameters +/// (cuda-macros `generate_device_function`, `has_codegen_generics`). It keeps +/// each monomorphization a separate item for the collector, a host-side +/// requirement rather than a device-code intent; as `noinline` it would keep +/// every call and stop const-generic arguments from folding into their callers. +/// The macro's attribute comes first, so rustc keeps it and reports an inline +/// attribute the user adds there as unused: on such functions a user's own +/// inline attribute never reached device code, before or after this mapping. +/// Unnamed instances never match (`opt_item_name`). The test keys on the +/// device symbol name plus own generics, which is enough while the macro +/// always writes its attribute first. Once #1472 lets a user's inline +/// attribute replace the macro's, this test must key on where the attribute +/// came from instead (`AttributeKind::Inline(_, span)` with +/// `span.from_expansion()`), or it would drop a user's own `#[inline(never)]`. +fn is_generic_device_collector_boundary<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> bool { + let def_id = instance.def_id(); + tcx.generics_of(def_id).own_requires_monomorphization() + && tcx + .opt_item_name(def_id) + .is_some_and(|name| reserved_oxide_symbols::is_device_symbol(name.as_str())) +} + /// Generates PTX for device functions using the cuda-oxide pipeline. /// /// This is the main entry point for device codegen. It bridges between @@ -997,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; @@ -1062,18 +1093,26 @@ pub fn generate_device_code<'tcx>( // 2. Sets up thread-local CompilerCtxt // 3. Runs our closure with access to stable() conversion // 4. Tears down the context and returns our result - // Pre-compute `#[inline(always)]` flags before entering the stable_mir - // context, since the query lives on `rustc_middle::TyCtxt` and is not - // exposed through stable_mir. Preserving this hint avoids making helper - // boundaries depend entirely on later optimizer heuristics. - let inline_always_flags: Vec = functions + // Pre-compute each function's `#[inline]` intent before entering the + // stable_mir context, since the query lives on `rustc_middle::TyCtxt` and + // is not exposed through stable_mir. Every variant is preserved (as rustc's + // LLVM backend does): helper boundaries, `#[inline(never)]` ones included, + // must not depend entirely on later optimizer heuristics. The one exception + // is the collector boundary `#[device]` adds to generic device functions. + let inline_intents: Vec> = functions .iter() .map(|func| { - let def_id = func.instance.def_id(); - matches!( - tcx.codegen_fn_attrs(def_id).inline, - rustc_hir::attrs::InlineAttr::Always | rustc_hir::attrs::InlineAttr::Force { .. } - ) + use dialect_mir::attributes::InlineIntentAttr; + use rustc_attr_ir::InlineAttr; + match tcx.codegen_fn_attrs(func.instance.def_id()).inline { + InlineAttr::None => None, + InlineAttr::Hint => Some(InlineIntentAttr::Hint), + InlineAttr::Always | InlineAttr::Force { .. } => Some(InlineIntentAttr::Always), + InlineAttr::Never if is_generic_device_collector_boundary(tcx, func.instance) => { + None + } + InlineAttr::Never => Some(InlineIntentAttr::Never), + } }) .collect(); let device_mono_reachability: Vec = functions @@ -1137,11 +1176,11 @@ pub fn generate_device_code<'tcx>( .iter() .zip(export_names.iter()) .zip(debug_scope_maps.iter()) - .zip(inline_always_flags.iter()) + .zip(inline_intents.iter()) .zip(device_mono_reachability.iter()) .filter_map( |( - (((func, (export_name, is_kernel)), debug_source_scopes), is_inline_always), + (((func, (export_name, is_kernel)), debug_source_scopes), inline), reachability, )| { // Use rustc_internal::stable() to convert the Instance. @@ -1172,7 +1211,7 @@ pub fn generate_device_code<'tcx>( export_name: export_name.clone(), debug_source_scopes: Some(debug_source_scopes.clone()), statement_debug_info, - is_inline_always: *is_inline_always, + inline: *inline, }) }, ) @@ -1215,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, @@ -1377,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 2c0b79d2e7..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,8 +4,8 @@ */ // @generated by cuda-intrinsics-gen 0.1.0. DO NOT EDIT. -// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 016c9d9b55b84d99d56fe2690dce4b3ba7ded5318890996a7459a142ddb5edf7. -// LLVM source: 16696adcd119e6ba9cc175207d984d7021211acb. +// Catalog schema/version: 46/0.1.0; intrinsic ABI: v1; catalog SHA-256: 63f8d327315cc3c3753b458364b3dafc962f63eb205b2eca77bff5f2e1f7faf4. +// LLVM source: aab11ee0ea5ea9a1d462ad0bc1fae9d9bb643019. //! Generated collector predicates for intrinsic placeholder functions. @@ -99,16 +99,21 @@ pub(crate) const GENERATED_INTRINSIC_CANONICAL_PATHS: &[&str] = &[ "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0893", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0922", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0328", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0335", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0329", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0330", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0331", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0336", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0332", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0337", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0333", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0338", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0334", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0339", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0928", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0929", @@ -201,12 +206,16 @@ pub(crate) const GENERATED_INTRINSIC_CANONICAL_PATHS: &[&str] = &[ "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0083", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0373", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0371", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0822", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0259", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0824", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0261", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0823", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0260", + "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0825", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0262", "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0372", @@ -243,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", @@ -357,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", @@ -366,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", @@ -521,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", @@ -1131,16 +1144,21 @@ pub(crate) const GENERATED_INTRINSIC_PUBLIC_PATHS: &[&str] = &[ "cuda_intrinsics::tma::cp_async_bulk_prefetch_tensor_gather4_2d_l2", "cuda_intrinsics::tma::cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint", "cuda_intrinsics::tma::cp_async_bulk_tensor_1d_g2s", + "cuda_intrinsics::tma::cp_async_bulk_tensor_1d_g2s_cta", "cuda_intrinsics::tma::cp_async_bulk_tensor_1d_s2g", "cuda_intrinsics::tma::cp_async_bulk_tensor_2d_g2s", + "cuda_intrinsics::tma::cp_async_bulk_tensor_2d_g2s_cta", "cuda_intrinsics::tma::cp_async_bulk_tensor_2d_g2s_multicast", "cuda_intrinsics::tma::cp_async_bulk_tensor_2d_g2s_multicast_cg2", "cuda_intrinsics::tma::cp_async_bulk_tensor_2d_s2g", "cuda_intrinsics::tma::cp_async_bulk_tensor_3d_g2s", + "cuda_intrinsics::tma::cp_async_bulk_tensor_3d_g2s_cta", "cuda_intrinsics::tma::cp_async_bulk_tensor_3d_s2g", "cuda_intrinsics::tma::cp_async_bulk_tensor_4d_g2s", + "cuda_intrinsics::tma::cp_async_bulk_tensor_4d_g2s_cta", "cuda_intrinsics::tma::cp_async_bulk_tensor_4d_s2g", "cuda_intrinsics::tma::cp_async_bulk_tensor_5d_g2s", + "cuda_intrinsics::tma::cp_async_bulk_tensor_5d_g2s_cta", "cuda_intrinsics::tma::cp_async_bulk_tensor_5d_s2g", "cuda_intrinsics::tma::cp_async_bulk_tensor_reduce_add_im2col_3d", "cuda_intrinsics::tma::cp_async_bulk_tensor_reduce_add_im2col_4d", @@ -1233,12 +1251,16 @@ pub(crate) const GENERATED_INTRINSIC_PUBLIC_PATHS: &[&str] = &[ "cuda_intrinsics::convert::cvt_rn_relu_f16x2_f32", "cuda_intrinsics::convert::cvt_rn_relu_satfinite_tf32_f32", "cuda_intrinsics::convert::cvt_rn_relu_tf32_f32", + "cuda_intrinsics::convert::cvt_rn_satfinite_e4m3x2_bf16x2", "cuda_intrinsics::convert::cvt_rn_satfinite_e4m3x2_f16x2", "cuda_intrinsics::convert::cvt_rn_satfinite_e4m3x2_f32", + "cuda_intrinsics::convert::cvt_rn_satfinite_e5m2x2_bf16x2", "cuda_intrinsics::convert::cvt_rn_satfinite_e5m2x2_f16x2", "cuda_intrinsics::convert::cvt_rn_satfinite_e5m2x2_f32", + "cuda_intrinsics::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2", "cuda_intrinsics::convert::cvt_rn_satfinite_relu_e4m3x2_f16x2", "cuda_intrinsics::convert::cvt_rn_satfinite_relu_e4m3x2_f32", + "cuda_intrinsics::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2", "cuda_intrinsics::convert::cvt_rn_satfinite_relu_e5m2x2_f16x2", "cuda_intrinsics::convert::cvt_rn_satfinite_relu_e5m2x2_f32", "cuda_intrinsics::convert::cvt_rn_satfinite_tf32_f32", @@ -1275,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", @@ -1389,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", @@ -1398,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", @@ -1553,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", @@ -2169,16 +2195,21 @@ pub(crate) fn is_generated_intrinsic_canonical_path(path: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0893" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0922" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0328" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1034" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0335" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0329" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1035" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0330" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0331" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0336" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0332" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1036" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0337" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0333" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1037" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0338" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0334" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1038" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0339" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0928" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0929" @@ -2271,12 +2302,16 @@ pub(crate) fn is_generated_intrinsic_canonical_path(path: &str) -> bool { | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0083" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0373" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0371" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1030" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0822" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0259" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1032" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0824" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0261" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1031" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0823" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0260" + | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i1033" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0825" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0262" | "cuda_intrinsics::__cuda_oxide_intrinsic_abi_v1::i0372" @@ -2313,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" @@ -2427,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" @@ -2436,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" @@ -2591,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" @@ -3208,16 +3247,21 @@ pub(crate) fn is_generated_intrinsic_compatibility_path(path: &str) -> bool { | "cuda_device::tma::cp_async_bulk_prefetch_tensor_gather4_2d_l2" | "cuda_device::tma::cp_async_bulk_prefetch_tensor_gather4_2d_l2_cache_hint" | "cuda_device::tma::cp_async_bulk_tensor_1d_g2s" + | "cuda_device::tma::cp_async_bulk_tensor_1d_g2s_cta" | "cuda_device::tma::cp_async_bulk_tensor_1d_s2g" | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s" + | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_cta" | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_multicast" | "cuda_device::tma::cp_async_bulk_tensor_2d_g2s_multicast_cg2" | "cuda_device::tma::cp_async_bulk_tensor_2d_s2g" | "cuda_device::tma::cp_async_bulk_tensor_3d_g2s" + | "cuda_device::tma::cp_async_bulk_tensor_3d_g2s_cta" | "cuda_device::tma::cp_async_bulk_tensor_3d_s2g" | "cuda_device::tma::cp_async_bulk_tensor_4d_g2s" + | "cuda_device::tma::cp_async_bulk_tensor_4d_g2s_cta" | "cuda_device::tma::cp_async_bulk_tensor_4d_s2g" | "cuda_device::tma::cp_async_bulk_tensor_5d_g2s" + | "cuda_device::tma::cp_async_bulk_tensor_5d_g2s_cta" | "cuda_device::tma::cp_async_bulk_tensor_5d_s2g" | "cuda_device::tma::cp_async_bulk_tensor_reduce_add_im2col_3d" | "cuda_device::tma::cp_async_bulk_tensor_reduce_add_im2col_4d" @@ -3310,12 +3354,16 @@ pub(crate) fn is_generated_intrinsic_compatibility_path(path: &str) -> bool { | "cuda_device::convert::cvt_rn_relu_f16x2_f32" | "cuda_device::convert::cvt_rn_relu_satfinite_tf32_f32" | "cuda_device::convert::cvt_rn_relu_tf32_f32" + | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_bf16x2" | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f16x2" | "cuda_device::convert::cvt_rn_satfinite_e4m3x2_f32" + | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_bf16x2" | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f16x2" | "cuda_device::convert::cvt_rn_satfinite_e5m2x2_f32" + | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_bf16x2" | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f16x2" | "cuda_device::convert::cvt_rn_satfinite_relu_e4m3x2_f32" + | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_bf16x2" | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f16x2" | "cuda_device::convert::cvt_rn_satfinite_relu_e5m2x2_f32" | "cuda_device::convert::cvt_rn_satfinite_tf32_f32" @@ -3352,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" @@ -3467,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" @@ -3476,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" @@ -3631,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 8deaf0cf0f..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 | //! @@ -305,6 +308,7 @@ // Import rustc internal crates extern crate rustc_abi; extern crate rustc_ast; +extern crate rustc_attr_ir; extern crate rustc_codegen_ssa; extern crate rustc_data_structures; extern crate rustc_driver; @@ -374,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, @@ -386,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>, @@ -401,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(), @@ -411,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() @@ -426,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('-', "_") } @@ -522,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); @@ -615,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..."); } @@ -658,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, @@ -728,7 +844,7 @@ impl CodegenBackend for CudaCodegenBackend { tcx.dcx().fatal(format!( "[rustc_codegen_cuda] Internal compiler error in \ device codegen: {msg}. This is a bug in cuda-oxide. \ - Please file at https://github.com/NVlabs/cuda-oxide/issues" + Please file at https://github.com/NVIDIA/cuda-rust/issues" )); } Ok(Ok(result)) => { @@ -772,6 +888,25 @@ impl CodegenBackend for CudaCodegenBackend { "[rustc_codegen_cuda] Device codegen did not produce an embeddable artifact", ); } + if !collection_result.tile_exports.is_empty() { + match write_tile_export_descriptions( + &device_config.output_dir, + &device_config.output_name, + &collection_result.tile_exports, + ) { + Ok(path) => { + if self.config.verbose { + eprintln!( + "[rustc_codegen_cuda] Tile export descriptions written: {}", + path.display() + ); + } + } + Err(e) => tcx.dcx().fatal(format!( + "[rustc_codegen_cuda] Failed to write tile export descriptions: {e}" + )), + } + } Some(result) } Ok(Err(e)) => { @@ -850,6 +985,25 @@ impl CodegenBackend for CudaCodegenBackend { } } +/// Write `.tileexports.json` next to the `.ll`/`.ptx`: the JSON +/// array of this crate's `#[tile_export]` descriptions. `cargo oxide tilelib` +/// reads it to build the Tile IR interface and implementation modules; the +/// device artifact itself carries nothing from it. +fn write_tile_export_descriptions( + output_dir: &Path, + output_name: &str, + tile_exports: &[oxide_tile_export::TileExportDesc], +) -> std::io::Result { + let path = output_dir.join(format!( + "{output_name}.{}", + oxide_tile_export::DESCRIPTIONS_SUFFIX + )); + let mut text = oxide_tile_export::descriptions_to_json_pretty(tile_exports); + text.push('\n'); + std::fs::write(&path, text)?; + Ok(path) +} + #[allow(clippy::too_many_arguments)] fn write_device_artifact_object( tcx: TyCtxt<'_>, @@ -1236,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..939d7756c3 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/tests/fixtures/primitives_api_compile/Cargo.lock @@ -0,0 +1,773 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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 new file mode 100644 index 0000000000..9f509220ab --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/advanced/images/mbarrier-thread-election.svg @@ -0,0 +1,750 @@ + + + + + +ManagedBarrier: one thread runs mbarrier.init, chosen by its block-linear id +All threads call init_by(count, id) and inval_by(id); one thread acts. Main chose it by threadIdx.x alone, which names one thread only in a 1D block; #1349 uses the block-linear id. +THE BLOCK (future_apis, test_managed_barrier_2d) +MAIN (threadIdx_x() == selector) +FIXED (block_linear_thread_id_here() == selector) + +block_dim = (32, 4, 1): 128 threads, one warp per row +y +threadIdx.x = 0 +31 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +0 +1 +2 +3 + +(0, 0, 0): linear id 0; the thread init() and inval() select + +x == 0 in rows 1 to 3: threadIdx.x == 0 as well, like (0, 0, 0) + +(31, 3, 0): linear id 31 + 32 * 3 = 127; no thread has x == 127 + +future_apis/src/main.rs: test_managed_barrier_2d +// ALL threads call init_by - only linear thread 127 runs mbarrier.init +let bar = unsafe { bar.init_by(BLOCK_2D_THREADS, BARRIER_2D_THREAD) }; +let token = bar.arrive(); +bar.wait(token); +let _dead = unsafe { bar.inval_by(BARRIER_2D_THREAD) }; +BLOCK_2D_THREADS = 128; BARRIER_2D_THREAD = 127 = thread (31, 3, 0) +As landed, from_static(), init(), init_by, inval() and inval_by are +all-threads calls; one thread, named by the selector, acts. The +selector must be the same value in every thread and below the thread +count: an id at or past it names nobody; nothing is initialized. + +main: barrier.rs, init_by and inval_by +if crate::thread::threadIdx_x() == init_thread { +the x coordinate alone: one thread only in a 1D block +selector 0 (what init() and inval() pass) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +four threads (x == 0 in each row) run mbarrier.init +on one word: undefined (PTX ISA), silent on sm_120 +selector 127 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +no thread has threadIdx.x == 127, so nothing runs +mbarrier.init; 128 arrive on an uninitialized word +result: the launch fails +Error: DriverError(719, "unspecified launch failure") +the same example source against main's cuda-device + +fixed: barrier.rs, init_by and inval_by +if crate::thread::block_linear_thread_id_here() + == init_thread { +id = x + 32 * (y + 4 * z): x fastest, then y, then z +selector 0 (what init() and inval() pass) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +one thread, (0, 0, 0); 1D blocks keep the same thread +selector 127 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +one thread, (31, 3, 0): 31 + 32 * 3 = 127 +result: 128 of 128 threads correct +All 128 threads wrote 42 in a (32, 4, 1) block +memcheck, synccheck, racecheck: 0 errors, 0 hazards + +thread::block_linear_thread_id_here() +the shared helper #1351 landed; thread_rank() uses it +too, so the PR's private copy of the formula is gone + + +THE PTX ELECTION (the predicate before mbarrier.init; tabs shown as spaces; bra.uni lines, labels and the count mov.b32 left out) + +main: test_managed_barrier (the 1D test; init(), so selector 0) +mov.u32 %r1, %tid.x; +setp.ne.b32 %p1, %r1, 0; +@%p1 bra $L__BB4_2; +mbarrier.init.shared.b64 [__shared_mem_e69151d2fb141f63_4], %r2; + +fixed: test_managed_barrier_2d (selector 127) +mov.u32 %r2, %tid.x; +mov.u32 %r3, %tid.y; +mov.u32 %r4, %tid.z; +mov.u32 %r5, %ntid.x; +mov.u32 %r6, %ntid.y; +mad.lo.s32 %r7, %r6, %r4, %r3; +// y + ntid.y * z +mad.lo.s32 %r8, %r7, %r5, %r2; +// x + ntid.x * (y + ntid.y * z) +// 7 lines left out: thread 0 stores all-ones into the barrier word, bar.sync +setp.ne.b32 %p2, %r8, 127; +@%p2 bra $L__BB2_4; +mbarrier.init.shared.b64 [__shared_mem_ddf9d77ea1d0101f_0], %r9; +What this proves: on main, %p1 reads one register, %tid.x. With the fix, %p2 reads five +registers through two mad and compares tid.x + ntid.x * (tid.y + ntid.y * tid.z) with the +selector, so mbarrier.init runs on at most one thread per block for any in-range selector +in any block shape. In both trees the same predicate guards mbarrier.inval later. On main +it admits four threads for selector 0 in a (32, 4, 1) block and no thread for selector 127. + 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 new file mode 100644 index 0000000000..a66298136b --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/advanced/images/shared-memory-alignment.svg @@ -0,0 +1,99 @@ + + + + + +Shared memory alignment: the request is checked, storage is declared at max(ALIGN, abi_align(T)) +Before, a shared array's PTX .align came from the request alone or from an exporter guess (8 for every aggregate); the element type's own alignment never entered, and nobody checked the request. +THE REQUEST (Rust) +ON MAIN (the request or an exporter guess, unchecked) +AFTER #1345 (verified request, typed storage) + +u64 at ALIGN 1 +static mut DATA: SharedArray<u64, 2, 1> +touched before a SharedArray<u8, 33> pad + +declared at the request +.shared .align 1 .b8 __shared_mem_<hash>_5[16] +st.shared.b64 [%rd4], %rd3 +DriverError(716, "misaligned address") + +raised to align_of::<u64>() +.shared .align 8 .b8 __shared_mem_<hash>_5[16] +values read back intact; memcheck and synccheck clean + +align(16) element, natural ALIGN +#[repr(C, align(16))] struct V4([f32; 4]); +static mut VS: SharedArray<V4, 4> +ALIGN 0; the raise lives on the Rust type + +declared at the exporter's guess +.shared .align 8 .b8 __shared_mem_<hash>_2[64] +st.shared.v4.b32 [%rd3], {%r2, %r4, %r3, %r1} +DriverError(716, "misaligned address") + +raised to the Rust ABI alignment +.shared .align 16 .b8 __shared_mem_<hash>_2[64] +mir_type_abi_align carries the repr(align(16)) raise + +dynamic pool at ALIGN 1 +DynamicSharedArray::<u64, 1>::get() +one pool per kernel, declared once for every use + +pool declared at the request +.extern .shared .align 1 .b8 __dynamic_smem_<hash>_<fn>[] +weaker than the st.shared.b64 through it + +pool at its strongest typed use +.extern .shared .align 8 .b8 __dynamic_smem_<hash>_<fn>[] +max over get, get_raw and offset uses, then the launch contract + +ALIGN 3 +static mut BAD: SharedArray<u64, 4, 3> +not a power of two; nobody checked it + +reached LLVM; no PTX, no source position +@__shared_mem_<hash>_0 = addrspace(3) global + [4 x i64] zeroinitializer, align 3 +llc: error: alignment is not a power of two + +rejected by the MIR verifier at the use site +["src/main.rs": line: 20, column: 13] +MirSharedAllocOp alignment must be a +non-zero power of two, found 3 + + + +Why main declared too little, and why it faulted +ALIGN > 0 reached the global unchecked. ALIGN 0 left the global with no alignment, and the exporter guessed one: +the integer width in bytes, 4 for f32, 8 for every other element, and 0 for bool (which llc rejected). +The element type's ABI alignment never entered, so the declared .align could be weaker than the kernel's own stores. +Shared globals are inserted at the module front, so the first array a kernel touches lands last in the PTX, right +behind the pad. ptxas packed DATA and VS at offsets their st.shared.b64 and st.shared.v4.b32 could not use, and +the launch failed with CUDA error 716, misaligned address (reproduced on sm_120a for three kernels). + +Where the alignment comes from now (mir-lower) +required_storage_align(ctx, mir_ty, llvm_ty) = + max(mir_type_abi_align(ctx, mir_ty), // Rust raise + llvm_type_size_align(ctx, llvm_ty).1) // LLVM layout +Both helpers already serve alloca, load, store and parameters. +The second sees LLVM leaves the MIR layout cannot, such as the +[16 x i64] TmaDescriptor lowers to: it still compiles, align 8. + +The contract +ALIGN is 0 (natural) or a non-zero power of two for SharedArray, and a non-zero power of two for DynamicSharedArray; anything else is a compile error that names the use site. +Storage is never declared below align_of::<T>(): a smaller request is raised silently, a larger one is honored. The dynamic pool takes its strongest typed use plus the launch contract. + 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 new file mode 100644 index 0000000000..f90ac2e9b2 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/advanced/images/tma-cta-local.svg @@ -0,0 +1,139 @@ + + + + + +CTA-local TMA loads: a shared pointer that is shared by type, and the .shared::cta copy form +The cluster form converts a generic *mut u8 to the cluster window; on sm_120a ptxas adds a runtime CTA check whose call drops setmaxnreg. The CTA form takes a SharedPtr<u8>. +TODAY (cp_async_bulk_tensor_2d_g2s, generic *mut u8) +THE TYPE (SharedPtr<T>, from #1341) +NEW (cp_async_bulk_tensor_2d_g2s_cta, SharedPtr<u8>) + +tma_copy_2d_test: the cluster form +static mut TILE: SharedArray<f32, TILE_SIZE, 128> + = SharedArray::UNINIT; +static mut BAR: Barrier = Barrier::UNINIT; + +// dst: a generic *mut u8, any CTA of the cluster +cp_async_bulk_tensor_2d_g2s( + &raw mut TILE as *mut u8, + tensor_map, tile_x, tile_y, + &raw mut BAR, +); +Unchanged by this PR: 39 in-tree callers keep this form and +its cluster contract; the PTX is byte for byte the same. + +tma_copy.ll: the generic pointer is converted +%v21 = addrspacecast ptr addrspace(3) @TILE to ptr +%v22 = addrspacecast ptr %v21 to ptr addrspace(7) +call void @llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d( + ptr addrspace(7) %v22, ptr addrspace(3) @BAR, ptr %v7, + i32 %v9, i32 %v10, i16 0, i64 0, i1 0, i1 0, i32 0) +The destination arrives as a generic pointer (the user's cast), +so the lowering can only name the cluster window, addrspace(7). +On sm_100a the cluster form is native. On sm_120a ptxas cannot +prove a non-multicast .shared::cluster destination local, so +it inserts a runtime CTA check with a syscall fallback, and +that call makes it ignore setmaxnreg (C7506): the #1238 kernel +ran with 24 registers instead of 168. + +SharedPtr<T>: addrspace(3) by type +let t = unsafe { SharedArray::shared_ptr(&raw mut TILE) }; +let d = DynamicSharedArray::<u8, 128>::shared_ptr_at(1024); +let p = unsafe { SharedPtr::from_raw(generic) }; // asserts +t.add(i) t.byte_add(b) t.cast::<u8>() // keep the type +t.as_ptr() // the generic *mut T for loads and stores +t.addr_u32() // the 32-bit .shared::cta address +repr(transparent) over *mut SharedSlot<T>; the importer types +that pointee as shared, next to SharedArray and Barrier, and +the newtype lowers as one scalar, like ThreadIndex. + +shared_memory_helper: across helper boundaries +struct Region<T> { base: SharedPtr<T>, offset: usize } +struct Pair { a: SharedPtr<f32>, b: SharedPtr<f32> } +#[inline(never)] +unsafe fn shared_ptr_probe<T>(region: Region<T>) -> u64 +#[inline(never)] +unsafe fn write_pair(pair: Pair, index: usize, value: f32) +#[inline(never)] +unsafe fn select_slot(pair: Pair, index: usize, second: bool) + -> SharedPtr<f32> + +the .ll: what each boundary keeps +define i64 @..shared_ptr_probe..( + ptr addrspace(3) %v0, i64 %v1) +define void @..write_pair( + ptr addrspace(3) %v0, ptr addrspace(3) %v1, i64, float) +define ptr addrspace(3) @..select_slot( + ptr addrspace(3) %v0, ptr addrspace(3) %v1, i64, i1) +Both routes; 0 cvta.to.shared in the PTX. A SharedPtr kernel +parameter, bare or in a struct or array, is a compile error. + +tma_copy_2d_cta_test: the CTA form (new) +// dst: SharedPtr<u8>, this CTA's shared memory by type +cp_async_bulk_tensor_2d_g2s_cta( + SharedArray::shared_ptr(&raw mut TILE).cast::<u8>(), + tensor_map, tile_x, tile_y, + &raw mut BAR, +); +cp_async_bulk_tensor_{1..5}d_g2s_cta, catalog records +i1034 to i1038, PTX 8.6, sm_90 and newer; tile mode only. + +LLVM NVPTX route: typed intrinsic (LLVM 22+) +call void @llvm.nvvm.cp.async.bulk.tensor.g2s.cta.tile.2d( + ptr addrspace(3) @TILE, ptr addrspace(3) @BAR, ptr %v7, + i32 %v9, i32 %v10, i64 0, i1 0) +No addrspacecast: the destination is already addrspace(3). + +libNVVM route: inline PTX, u32 addresses +asm sideeffect + "cp.async.bulk.tensor.2d.shared::cta.global.tile + .mbarrier::complete_tx::bytes [$0], [$2, {$3, $4}], [$1];", + "r,r,l,r,r,~{memory}" + +refused at compile time +a SharedPtr kernel parameter, bare or inside a struct or array +a barrier passed straight from map_shared_rank_mut (a located +error that names the cluster form to use instead) +the caller's contract: from_raw on a cluster pointer; the raw +cuda_intrinsics layer keeps *mut u8 and converts with +cvta.to.shared, and its safety text says CTA-local + + +THE PTX AND WHAT PTXAS MAKES OF IT ON sm_120a (tma_copy.ptx built for sm_90a; mnemonic column aligned; the copy line wrapped before .mbarrier; register counts from ptxas 13.4) + +cluster form: tma_copy_2d_test (unchanged) +mov.b64 %rd15, TILE; +cvta.shared.u64 %rd16, %rd15; +cvta.to.shared::cluster.u64 %rd17, %rd16; +cp.async.bulk.tensor.2d.shared::cluster.global.tile + .mbarrier::complete_tx::bytes [%rd17], [%rd2, {%r3, %r4}], [BAR]; +// ptxas sm_120a: 24 registers, 8-byte stack; with setmaxnreg: C7506 ignored +// SASS: CALL.ABS.NOINC, then UTMALDG.2D; no USETMAXREG + +CTA form: tma_copy_2d_cta_test +// no address conversion: the shared symbols are the operands +cp.async.bulk.tensor.2d.shared::cta.global.tile + .mbarrier::complete_tx::bytes [TILE], [%rd2, {%r3, %r4}], [BAR]; +// ptxas sm_120a: 10 registers, no stack; with setmaxnreg: 168 registers +// no C7506; SASS: USETMAXREG.DEALLOC.CTAPOOL 0x28, UTMALDG.2D, +// USETMAXREG.TRY_ALLOC.CTAPOOL 0xe8; no CALL +What this proves: the cluster form converts the generic destination twice (to generic, then to the +cluster window) and ptxas on sm_120a adds the runtime CTA check whose call drops setmaxnreg. The CTA +form names the shared symbols directly and keeps register reallocation. On the 5090, tma_copy copies +4096 of 4096 values through the new kernel next to the unchanged ones, through both routes, with +memcheck and synccheck clean; ranks 1 to 5 are probed and assembled, and the cluster PTX is unchanged. + diff --git a/cuda-oxide/cuda-oxide-book/advanced/matrix-multiply-accelerators.md b/cuda-oxide/cuda-oxide-book/advanced/matrix-multiply-accelerators.md index 57ec22630c..c0dfede1e1 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/matrix-multiply-accelerators.md +++ b/cuda-oxide/cuda-oxide-book/advanced/matrix-multiply-accelerators.md @@ -376,7 +376,7 @@ for tiles, `ManagedBarrier` for synchronization, and the MMA APIs for compute. The [`gemm_sol` -example](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol) +example](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol) is the worked-out reference. Its 4-stage `cta_group::2` pipeline reaches **868 TFLOPS at 4096³ — 57.8 % of `cublasLtMatmul` SoL — on B200 (148 SMs)**. Absolute throughput scales with SM count on smaller Blackwell datacenter diff --git a/cuda-oxide/cuda-oxide-book/advanced/shared-memory-and-synchronization.md b/cuda-oxide/cuda-oxide-book/advanced/shared-memory-and-synchronization.md index c068c8af4f..5871909e1a 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/shared-memory-and-synchronization.md +++ b/cuda-oxide/cuda-oxide-book/advanced/shared-memory-and-synchronization.md @@ -123,7 +123,7 @@ undefined until threads write to them. | :-------- | :-------------------------------------------------------------------------------- | | `T` | Element type (`f32`, `u32`, etc.) | | `N` | Number of elements (fixed at compile time) | -| `ALIGN` | Byte alignment (default 0 = natural alignment; use 128 for TMA destinations) | +| `ALIGN` | Byte alignment: 0 (natural) or a power of two; use 128 for TMA destinations | :::{tip} `SharedArray` is `!Sync` by design — it wraps `UnsafeCell`, which prevents @@ -132,6 +132,33 @@ shared memory is inherently mutable by all threads in the block, and the programmer is responsible for synchronization. ::: +### Alignment rules + +`ALIGN` is `0` (natural alignment) or a non-zero power of two. Any other +value is a compile error that names the use site. + +Storage is never declared below `align_of::()`. A smaller request is +raised silently, the way `#[repr(align(N))]` raises a struct; a larger +request is honored. `SharedArray` is therefore declared at 8 +bytes, and a `#[repr(C, align(16))]` element at natural alignment is +declared at 16. Without the raise, a `u64` array declared at 1 byte could be +packed right behind a byte array and fault with `misaligned address`. + +Aggregate, nested-array and `f16` elements at natural alignment declare +their natural alignment in PTX (`.align 4` for `[f32; N]`, `.align 2` for +`f16`). Older releases printed a blanket `.align 8` for them. + +```{figure} images/shared-memory-alignment.svg +:align: center +:width: 100% + +Shared-memory storage alignment. The verifier rejects an `ALIGN` that is not a non-zero power +of two at the use site, and lowering declares every static array and the dynamic pool at the +larger of the request and the element's ABI alignment. `SharedArray` is declared +`.align 8` and a `#[repr(C, align(16))]` element at natural alignment `.align 16`, where older +releases declared `.align 1` and `.align 8` and the kernel could fault with `misaligned address`. +``` + ### The API surface ```rust @@ -140,6 +167,7 @@ impl SharedArray { pub const fn len() -> usize; pub fn as_ptr(&self) -> *const T; pub fn as_mut_ptr(&mut self) -> *mut T; + pub unsafe fn as_raw_mut_ptr(shared: *mut Self) -> *mut T; } // Indexing (unsafe access via static mut) @@ -233,6 +261,78 @@ total does not exceed `shared_mem_bytes` — there are no runtime guards. (Hopper+), use `DynamicSharedArray` to get the required 128-byte alignment. The alignment is encoded in the PTX `.align` directive. +`ALIGN` must be a non-zero power of two; `0` is a compile error that names +the use site. A kernel has one dynamic pool, and the compiler declares it at +the strongest typed requirement over the kernel's `get`, `get_raw` and +`offset` uses, each taken as the larger of its `ALIGN` and +`align_of::()`, combined with the launch contract's alignment. +`DynamicSharedArray::::get()` alone declares the pool at 8 bytes. +The declaration aligns the pool base only; keeping an `offset(byte_offset)` +result aligned for `T` is still the caller's job. + +### Typed shared pointers: `SharedPtr` + +`DynamicSharedArray::get()` and `SharedArray::as_mut_ptr()` return `*mut T`, +which in device code is a *generic* pointer: it may address global, local or +shared memory, including another block's shared memory in a cluster. Once +such a pointer is passed to a helper function, the compiler no longer knows it +points into this block's shared memory. + +`SharedPtr` is a raw pointer into the issuing block's shared memory +(`.shared::cta`) that carries that fact in its type. It lowers to an +`addrspace(3)` pointer wherever it goes, including helper parameters and +return values, struct fields and generic code, so instructions that take it +use their shared-memory forms without converting a generic address. + +```rust +use cuda_device::{DynamicSharedArray, SharedArray, SharedPtr, kernel, thread}; + +#[inline(never)] +unsafe fn fill(tile: SharedPtr, i: usize, value: f32) { + // Still a shared-memory pointer here, although `fill` is not inlined. + unsafe { tile.add(i).as_ptr().write(value) }; +} + +static mut TILE: SharedArray = SharedArray::UNINIT; + +#[kernel] +pub fn fill_tile() { + // SAFETY: TILE is a `static mut SharedArray` of this kernel. + let tile: SharedPtr = unsafe { SharedArray::shared_ptr(&raw mut TILE) }; + // The same type for dynamic shared memory, 1024 bytes into the pool. + let pool: SharedPtr = DynamicSharedArray::::shared_ptr_at(1024); + // SAFETY: the launch supplies more than 1024 dynamic shared bytes, and + // each thread writes its own slot of TILE. + unsafe { + fill(tile, thread::threadIdx_x() as usize, 1.0); + pool.as_ptr().write(0); + } +} +``` + +| Operation | Result | +|:---------------------------------------|:----------------------------------------------------| +| `unsafe SharedArray::shared_ptr(&raw mut S)` | `SharedPtr` to the first element of `S`, a `static mut` of this kernel | +| `DynamicSharedArray::::shared_ptr()` / `shared_ptr_at(bytes)` | Like `get()` / `offset(bytes)` | +| `unsafe SharedPtr::from_raw(p)` | Asserts that generic `p` points into this block's shared memory | +| `add`, `offset`, `byte_add`, `cast` | Pointer arithmetic that keeps the type | +| `as_ptr()` | The generic `*mut T`, for loads, stores and raw-pointer APIs | +| `addr_u32()` | The 32-bit `.shared::cta` address, for `_shared_u32` forms | + +`SharedPtr` cannot be a kernel parameter, bare or inside a struct or array: +the host has no shared-memory address to pass, and the compiler refuses the +signature. Pointers from `cluster::map_shared_rank` may name another block's +shared memory and are not `SharedPtr`s. + +Inside aggregates the usual shared-pointer limits apply. A union or a +`MaybeUninit>` fails closed under the union storage rules, and a +`repr(packed)` struct that holds one follows the packed-carrier limits; both +are listed in [Supported Features](../appendix/supported-features.md). On the +libNVVM route a shared pointer is 4 bytes wide, so a struct that mixes a +`SharedPtr` with a narrower field, such as `{ p: SharedPtr, x: u32 }`, is +not yet supported in arrays or in shared memory. That gap predates +`SharedPtr` and applies to `*mut Barrier` fields in the same way. + --- ## Synchronization: sync_threads() @@ -272,6 +372,105 @@ that not all threads take. Every thread in the block must reach the same control flow, restructure so the barrier is outside the branch. ::: +(compiler-ordering-around-the-barrier)= + +### Compiler ordering around the barrier + +The hardware guarantee above covers visibility between threads. A second +question is whether the compiler keeps your own accesses on the side of the +barrier where you wrote them. Both backends emit the barrier as an operation +that the compiler treats as reading and writing all memory: + +- The LLVM NVPTX route calls `llvm.nvvm.barrier.cta.sync.aligned.all`. The + call is `convergent` and has no `memory(...)` attribute. +- The libNVVM route emits `bar.sync 0;` as `asm sideeffect` with a `~{memory}` + clobber. + +| Access form | Kept on its side of `sync_threads()` | Notes | +| :------------------------------------------------------- | :----------------------------------- | :---- | +| `SharedArray` indexing, `S[i]` | Yes | A plain load or store in address space 3. | +| `SharedArray::as_raw_mut_ptr(&raw mut S)` | Yes | The same instructions as indexing. Prefer it to `as_mut_ptr()` and `as_ptr()`. Those methods borrow `S`, and the `static_mut_refs` lint rejects a borrow of a `static mut` under edition 2024. | +| `DynamicSharedArray::::get()` and `offset()` | Yes | A plain load or store through the returned pointer. | +| `ptx_asm!` with no options | Yes | The call is `asm sideeffect` with no memory attribute. The PTX keeps it in order against the barrier, plain loads and stores, and other side-effecting asm. ptxas may drop a snippet with an unused result from the SASS. | +| `ptx_asm!` with `clobber("memory")` | Yes | Adds `~{memory}` to the constraint string. The order matches the row above in the current lowering. Without the clobber that order comes from the `sideeffect` flag alone. That flag does not tell the compiler that the snippet touches memory. Add it to every snippet that loads or stores through an address. | +| `ptx_asm!` with `options(register_only)` | No | The call has no `sideeffect` flag. The compiler may move it across memory operations and the barrier, and deletes it when its result is unused. Use it for snippets that touch registers only. | +| `cp_async_*` copies, `cp_async_commit_group()`, `cp_async_wait_group(n)` | Yes, but see below | Issue, commit and wait stay in order. The barrier does not wait for the copy to land. | +| A `&T` or `&[T]` to slots another thread writes, live across the barrier | No | rustc merges reads through it before either backend runs. See below. | +| A `&UnsafeCell` to such a slot | Yes | A plain load that is re-read after each barrier. Two adjacent reads with nothing between them may still become one load. | +| A reference to a `BlockAtomic*` type | Yes | Each `load()` is side-effecting asm with a memory clobber. The compiler keeps every one. | + +The indexing, `as_raw_mut_ptr`, `DynamicSharedArray` and both side-effecting +`ptx_asm!` rows were checked on both backends with a producer and consumer +kernel: store, barrier, load, barrier, store, barrier, load. The PTX keeps that +order for each of them. It keeps the same order when the loads sit in a loop +whose trip count is a kernel argument. A run on an RTX 5060 Laptop (sm_120) +read the expected neighbour values. A `register_only` `ld.shared` whose result +is unused leaves no instruction in the PTX. The same snippet with no options +stays. The `cp.async` and reference rows have their own tests below. + +The table covers compiler ordering only. The rule under "When to sync" still +decides where the barriers belong. + +#### Raw pointers across the barrier + +Each thread derives its pointer from the raw address of the allocation and +creates no reference to the whole array. A thread writes its own slot +before the barrier and reads a neighbour's slot after it: + +```rust +static mut TILE: SharedArray = SharedArray::UNINIT; + +let tid = thread::threadIdx_x() as usize; +let tile = unsafe { SharedArray::as_raw_mut_ptr(&raw mut TILE) }; +unsafe { tile.add(tid).write(tid as u32 + 1) }; // this thread's slot +thread::sync_threads(); +let next = unsafe { tile.add((tid + 1) % 128).read() }; // a neighbour's slot +``` + +A kernel built from this snippet read the neighbour's value in all 128 threads +on both backends. + +#### `cp.async` needs a wait before the barrier + +`bar.sync` does not wait for a pending `cp.async`. Each thread waits for its +own copies and then every thread meets at the barrier: + +```rust +unsafe { + cp_async_ca_4(dst, src); + cp_async_commit_group(); + cp_async_wait_group(0); // this thread's copies have landed +} +thread::sync_threads(); // other threads may now read `dst` +``` + +Reading the destination before the copy completes is undefined behaviour. A +test kernel that drops `cp_async_wait_group(0)` and reads a neighbour's slot +after the barrier can return the stale value: in a plain run on sm_120 all 128 +threads did, and under compute-sanitizer every copy had landed first. The same +kernel with the wait returned the copied value in all 128 threads. + +#### Share a slot through a reference only with `UnsafeCell` + +When `T` has no interior mutability, a `&T` into shared memory tells rustc +that the target does not change while the reference is live. rustc merges two +reads through it into one load before either backend runs. In a test kernel +each thread read its neighbour's slot through a `&u32`, passed two barriers +while the neighbour wrote a new value, and read again. The second read +returned the old value in all 128 threads on both backends. A `&[T]` built +over the tile with `from_raw_parts` is the same case: two indexed reads across +the barrier compile to one load. Keeping such a reference live across a barrier +while another thread writes the target is undefined behaviour. So is a `&mut T` +to such a slot: the other thread's write breaks its exclusive access, even +though no pass merges reads through it today. + +A `&UnsafeCell` is re-read after each barrier, and every `BlockAtomic*` +load is an instruction the compiler keeps. The same kernel with either form +read the new value in all 128 threads. +Share a slot through one of these, or read through a raw pointer after each +barrier. The tip under "Declaration rules" gives the same reason for +`SharedArray` wrapping `UnsafeCell`. + --- ## Shared memory vs. other approaches diff --git a/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md b/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md index a4b9da986a..f2397747bf 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md +++ b/cuda-oxide/cuda-oxide-book/advanced/tensor-memory-accelerator.md @@ -119,6 +119,42 @@ The coordinates (`tile_x`, `tile_y`) specify which tile to copy, in units of the box dimensions encoded in the descriptor. The barrier pointer tells the TMA engine where to signal completion. +These functions address cluster shared memory: the destination may belong to +any CTA of the cluster, so the compiler converts it to the cluster window. +When the destination is the issuing CTA's own shared memory, +`cp_async_bulk_tensor_*_g2s_cta` emits the `.shared::cta` destination form +(PTX 8.6, sm_90+) without that conversion. Its destination is a +[`SharedPtr`](shared-memory-and-synchronization.md), +a CTA-shared pointer by type, so the locality holds wherever the pointer came +from, including through helper functions: + +```rust +// SAFETY: TILE and BAR are `static mut` shared-memory objects of this kernel, +// and the copy completes on BAR before any thread reads TILE. +unsafe { + cp_async_bulk_tensor_2d_g2s_cta( + SharedArray::shared_ptr(&raw mut TILE).cast::(), + tensor_map, tile_x, tile_y, + &raw mut BAR, + ); +} +``` + +On SM12x this form matters for performance: ptxas implements a non-multicast +`.shared::cluster` TMA load it cannot prove CTA-local with a runtime check and +a system call, which also makes it ignore `setmaxnreg`. + +```{figure} images/tma-cta-local.svg +:align: center +:width: 100% + +The two G2S destination forms. Left, the cluster form converts a generic `*mut u8` to the cluster +window, and on sm_120a ptxas adds a runtime CTA check whose call makes it ignore `setmaxnreg`. +Middle, `SharedPtr` keeps the shared address space in its type across helper parameters, struct +fields and return values. Right, `cp_async_bulk_tensor_2d_g2s_cta` takes a `SharedPtr` and emits +the `.shared::cta` form with no address conversion, so the kernel keeps `setmaxnreg` and its 168 registers. +``` + ### Only one thread issues the copy TMA copies are **not** collective operations across the block. Exactly @@ -188,6 +224,18 @@ The lifecycle is: `Uninit` → `init()` → `Ready` → `arrive()` / `wait()` `inval()` → `Invalidated`. The type system prevents calling `wait` on an uninitialized barrier or `arrive` on an invalidated one. +```{figure} images/mbarrier-thread-election.svg +:align: center +:width: 100% + +Which thread runs `mbarrier.init`. Every thread calls `init_by(count, id)` and `inval_by(id)`, +and the thread whose block-linear id (`x + blockDim.x * (y + blockDim.y * z)`, the numbering +`thread_rank()` uses) equals `id` runs the instruction; `init()` and `inval()` pass 0, thread +(0, 0, 0). Comparing `threadIdx.x` alone named one thread only in a 1D block: in a (32, 4, 1) +block selector 0 ran the instruction on four threads and selector 127 on none, so the launch +failed with error 719. +``` + ### Key barrier rules | Rule | Why | diff --git a/cuda-oxide/cuda-oxide-book/advanced/warp-level-programming.md b/cuda-oxide/cuda-oxide-book/advanced/warp-level-programming.md index 7e50daaec5..4df7a75d2a 100644 --- a/cuda-oxide/cuda-oxide-book/advanced/warp-level-programming.md +++ b/cuda-oxide/cuda-oxide-book/advanced/warp-level-programming.md @@ -28,9 +28,10 @@ let lane = warp::lane_id(); // 0..31, hardware register %laneid let warp = warp::warp_id(); // threadIdx.x / 32 ``` -`warp_id()` is derived from `threadIdx.x / 32`. For multi-dimensional -blocks, this only accounts for the x dimension — which is usually fine, -since most kernels that care about lane identity use 1D blocks. +`warp_id()` is derived from `threadIdx.x / 32`. In a 2D or 3D block that +value does not name a hardware warp: warps are formed from consecutive +block-linear thread ids, `x + blockDim.x * (y + blockDim.y * z)`. Use +`this_thread_block().thread_rank() / 32` there. --- 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 9bf9e49ad3..cc90fe89a5 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md +++ b/cuda-oxide/cuda-oxide-book/appendix/api-quick-reference.md @@ -55,7 +55,8 @@ fn helper(x: f32) -> f32 { x * x } | `#[readonly]` | Mark as read-only | Use these annotations only on an explicit counted `while` loop or a range -`for` loop (`for i in a..b`) inside a `#[kernel]` or `#[device]` function. +`for` loop (`for i in a..b` or `(a..b).rev()`, over any integer counter +including `isize`) inside a `#[kernel]` or `#[device]` function. Other iterator loops are not yet recognized by the unroll pass. Nested loops and multiple `continue` paths are supported. Full `#[unroll]` preserves `break` paths and multiple exit targets. @@ -72,7 +73,7 @@ and 65,536 cloned operations. Larger requests warn and are not unrolled. ```rust use cuda_device::{gpu_assert, gpu_printf, ptx_asm}; -gpu_printf!("thread %d: val = %f\n", idx as i32, val as f64); +gpu_printf!("thread {}: val = {}\n", idx, val); gpu_assert!(val.is_finite()); gpu_assert!(val >= 0.0, "expected non-negative value"); @@ -193,7 +194,7 @@ Generic policy expressions require `#![feature(generic_const_exprs)]`. currently remain literal. See the -[`policy_config`](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config) +[`policy_config`](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config) example for two concrete policies that generate independent PTX specializations and policy-specific prepared launches. @@ -311,6 +312,7 @@ pub fn dynamic(data: &[f32]) { | `SharedArray` | Compile-time sized, block-scoped shared memory | | `SharedArray` | With 128-byte alignment (required for TMA destinations) | | `DynamicSharedArray` | Runtime-sized shared memory (set via `LaunchConfig`) | +| `SharedPtr` | Pointer into this block's shared memory, `addrspace(3)` by type | Both are `!Sync` — concurrent access requires explicit barriers. @@ -449,6 +451,7 @@ cp_async_bulk_commit_group(); | Function | Description | |:----------------------------------------------|:-------------------------------------| | `cp_async_bulk_tensor_{1..5}d_g2s(...)` | Global → shared async bulk copy | +| `cp_async_bulk_tensor_{1..5}d_g2s_cta(...)` | Global → this CTA's `SharedPtr` | | `cp_async_bulk_tensor_{1..5}d_s2g(...)` | Shared → global async bulk copy | | `cp_async_bulk_tensor_2d_g2s_multicast(...)` | Multicast to all CTAs in cluster | | `cp_async_bulk_commit_group()` | Commit outstanding copies | @@ -648,7 +651,7 @@ debug::prof_trigger::<7>(); // Nsight profiler trigger | `thread` | Thread/block IDs, `index_1d`, `sync_threads` | All | | `config` | Compile-time policies, shapes, tiles, atoms, layouts, memory spaces, and scopes | All | | `disjoint` | `DisjointSlice` — typed writes completed by a launch proof | All | -| `shared` | `SharedArray`, `DynamicSharedArray` | All | +| `shared` | `SharedArray`, `DynamicSharedArray`, `SharedPtr` | All | | `warp` | Shuffle, vote, match, lane/warp ID | All | | `atomic` | Scoped atomics (device/block/system) | sm_70+ | | `debug` | `clock64`, `trap`, `breakpoint`, `gpu_printf!` | All | diff --git a/cuda-oxide/cuda-oxide-book/appendix/building-from-source.md b/cuda-oxide/cuda-oxide-book/appendix/building-from-source.md index 598843eb77..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-08-28` (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 | @@ -19,29 +19,28 @@ from a fresh checkout. If you just want to run an example, the ## Clone the repository ```bash -git clone https://github.com/NVlabs/cuda-oxide.git -cd cuda-oxide +git clone https://github.com/NVIDIA/cuda-rust.git +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-08-28" +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-08-28 -rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain nightly-2026-08-28 +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-08-28`) 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: @@ -134,28 +133,33 @@ sudo apt install clang-21 # or libclang-common-21-dev ## Build the workspace -The main workspace contains the user-facing 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 comes -from crates.io: +The SIMT workspace lives under `cuda-oxide/`. It contains the user-facing +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 SIMT toolchain pin and the +`cargo oxide` workspace alias apply: ```bash +cd cuda-oxide 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-08-28 install --git https://github.com/NVlabs/cuda-oxide.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 @@ -163,6 +167,8 @@ dylib. ## Verify the installation +From `cuda-oxide/`: + ```bash # Check all prerequisites cargo oxide doctor @@ -179,6 +185,8 @@ and prints a success message. ## Common commands +From `cuda-oxide/`: + ```bash # Build and run an example cargo oxide run @@ -208,6 +216,7 @@ The documentation lives in `cuda-oxide-book/` and uses Sphinx with MyST Markdown. To build and serve locally: ```bash +# From the repository root: cd cuda-oxide/cuda-oxide-book make setup # creates venv, installs dependencies source .venv/bin/activate @@ -219,14 +228,15 @@ make livehtml # starts dev server on http://localhost:8000 Standard `cargo doc` works for the workspace crates: ```bash +# From the repository root: cd cuda-oxide cargo doc --no-deps --open ``` This generates rustdoc for `cuda-device`, `cuda-host`, `cuda-macros`, and all other workspace members. The codegen backend is excluded since it is not a -workspace member, and the shared host runtime (`cuda-core`, `cuda-async`) is -documented from cutile-rs. +workspace member. The shared host runtime (`cuda-core`, `cuda-async`) lives in +the workspace at the repository root; run `cargo doc` there to document it. ## Workspace structure @@ -239,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 @@ -277,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-08-28`. + `rustup component add rustc-dev --toolchain 1.99.1`. CUDA driver version mismatch : The toolkit version (compile-time) and driver version (runtime) must be @@ -286,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/ecosystem.md b/cuda-oxide/cuda-oxide-book/appendix/ecosystem.md index 050c093b36..238153342a 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/ecosystem.md +++ b/cuda-oxide/cuda-oxide-book/appendix/ecosystem.md @@ -90,5 +90,5 @@ If you're building a Rust + GPU project — or evaluating which of the above fits your needs — we're happy to compare notes. Join the community on [Discord](https://discord.gg/ZUEr4AhH5C) for questions, design discussions, and announcements, or reach the team via -[GitHub Discussions](https://github.com/NVlabs/cuda-oxide/discussions) +[GitHub Discussions](https://github.com/NVIDIA/cuda-rust/discussions) or by opening an issue on the repository. diff --git a/cuda-oxide/cuda-oxide-book/appendix/glossary.md b/cuda-oxide/cuda-oxide-book/appendix/glossary.md index a6a1a3a757..8c66a44cec 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/glossary.md +++ b/cuda-oxide/cuda-oxide-book/appendix/glossary.md @@ -38,8 +38,8 @@ alongside the normal host binary. ## `cuda-async` -The async execution layer, a crate shared with cutile-rs and published from -NVlabs/cutile-rs. The cuda-oxide model lives under `cuda_async::simt`: +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`: `DeviceOperation` (lazy GPU work description), `DeviceFuture` (stream-bound execution), and `DeviceBox` (device-owned memory). Compose work with `zip!`, `and_then`, and `value()`. @@ -54,8 +54,8 @@ 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 published from NVlabs/cutile-rs; -the cuda-oxide surface is re-exported at the crate root except +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`. ## `DeviceOperation` diff --git a/cuda-oxide/cuda-oxide-book/appendix/supported-features.md b/cuda-oxide/cuda-oxide-book/appendix/supported-features.md index 2fc179653f..71a5df7b38 100644 --- a/cuda-oxide/cuda-oxide-book/appendix/supported-features.md +++ b/cuda-oxide/cuda-oxide-book/appendix/supported-features.md @@ -106,6 +106,24 @@ thin-pointer-only device-global union initializers may preserve one full-width relocation at byte zero; mixed pointer/integer device-global initializers remain rejected. +A union field that holds a shared-memory pointer, directly or through structs +and bounded arrays, uses generic-pointer storage, so the union's byte image is +the same on every target and its integer view equals `ptr as usize`. Use +`cvta_generic_to_shared_offset` or `cvta_generic_to_shared_u32` when the +shared-window offset (the pointer's position inside the block's shared memory) +is wanted instead. A field whose fixed arrays expand to more than 16 +shared-memory pointers is rejected; pointers held directly or through structs +do not count toward that bound. Taking the in-place address of such a +field (`u.ptr = p`, `&u.ptr`, `&mut u.ptr`) is refused with a diagnostic; build +the union by value and read its fields by value. + +`SharedPtr`, the typed shared pointer, follows the same limits: a union or a +`MaybeUninit>` fails closed under these union storage rules, a +`repr(packed)` struct that holds one follows the packed-carrier limits above, a +kernel parameter that holds one bare or inside a struct or array is refused, +and on the libNVVM route a struct that mixes one with a narrower field is not +yet supported in arrays or in shared memory. + Enum constants preserve payload relocations to device statics, including non-zero byte addends. This includes niche-encoded `Option<&T>` and direct-tagged enum layouts, both for direct thin-reference payloads and for @@ -131,7 +149,7 @@ Anonymous promoted allocations remain unsupported. | For Loops (range, iterator, enumerate) | **Full** | Full iterator desugaring: range-based, `slice.iter()`, `enumerate()`, nested loops, `break`, `continue`. | | While Loops / If-Else | **Full** | Baseline control flow fully supported. | | Break and Continue | **Full** | `break` and `continue` in for/while loops, including early exit. | -| Loop Unroll Annotations | **Partial** | `#[unroll]` and `#[unroll(N)]` request unrolling of explicit counted `while` loops. Nested loops and multiple `continue` paths work; full unrolling preserves `break` paths and multiple exit targets, while partial unrolling requires a positive-step `<`/`<=` loop with an invariant limit and only the normal header exit. Requests are capped at 1,024 copies, 8,192 cloned blocks, and 65,536 cloned operations. | +| Loop Unroll Annotations | **Partial** | `#[unroll]` and `#[unroll(N)]` request unrolling of explicit counted `while` loops and range `for` loops (`for i in a..b` and `(a..b).rev()`) over any integer counter, including `isize`. Other iterator loops, such as `a..=b` or `.step_by(k)`, are not yet recognized. Nested loops and multiple `continue` paths work; full unrolling preserves `break` paths and multiple exit targets, while partial unrolling requires a positive-step `<`/`<=` loop with an invariant limit and only the normal header exit. Requests are capped at 1,024 copies, 8,192 cloned blocks, and 65,536 cloned operations. | | Monomorphization-Dead Branches | **Partial** | Branches that become dead after generic specialization (e.g. the const-false arm of `if M::ENABLED`) are ignored by symbol collection, panic checks, and pointer address-space inference, so panic-only hooks in dead arms compile. Only switches rustc itself folds are pruned: a constant discriminant operand or a direct single-assignment constant. Multi-step constant copy chains keep both arms, matching rustc's host monomorphization; this is deliberate, not a general constant-propagation pass. | ## Compiler: Arithmetic and Casting @@ -149,7 +167,7 @@ Anonymous promoted allocations remain unsupported. | Bi-directional LTOIR Support | **Full** | Rust kernels call CUDA C++ device functions **and** C++ calls Rust device functions. Via NVVM IR → libNVVM → LTOIR → nvJitLink. | | Device FFI (`extern "C"`) | **Full** | `#[device] extern "C" { fn ... }` declarations for external LTOIR functions. CUB/CCCL integration demonstrated. | | MathDx FFI (cuFFTDx / cuBLASDx) | **Full** | cuFFTDx (8/16/32-point thread-level FFT), cuBLASDx (32x32x32 block-level GEMM) via LTOIR. | -| Tile interop | **Partial** | Inter-kernel interop works today: a [cutile-rs Tile kernel](https://github.com/NVlabs/cutile-rs) and a cuda-oxide SIMT PTX kernel can run in one host process on the same CUDA stream over shared device tensors. Intra-kernel Tile interop is work in progress and tracked in [#96](https://github.com/NVlabs/cuda-oxide/issues/96). | +| Tile interop | **Partial** | Inter-kernel interop works today: a [cutile-rs Tile kernel](https://github.com/NVIDIA/cuda-rust/tree/main/cutile-rs) and a cuda-oxide SIMT PTX kernel can run in one host process on the same CUDA stream over shared device tensors. Intra-kernel interop is tracked in [#96](https://github.com/NVlabs/cuda-oxide/issues/96). | | Cross-Crate Kernels | **Full** | Kernels and device functions defined in library crates with monomorphization at the binary crate use site. | ## Compiler: Functions @@ -157,7 +175,7 @@ Anonymous promoted allocations remain unsupported. | Feature | Status | Description | |:--------|:-------|:------------| | `#[kernel]` Attribute | **Full** | Marks functions as GPU kernel entry points (`ptx_kernel` calling convention). Multiple kernels per file. | -| `#[device]` Helper Functions | **Full** | Device-side helper functions callable from kernels. `#[inline(always)]` is preserved as the LLVM `alwaysinline` attribute (emitted alongside the convergent group and any `!dbg` scope), so `opt` honors the inline intent. | +| `#[device]` Helper Functions | **Full** | Device-side helper functions callable from kernels. The `#[inline]` family reaches LLVM as the matching function attribute, as in rustc's own backend: `#[inline]` becomes `inlinehint`, `#[inline(always)]` becomes `alwaysinline` and `#[inline(never)]` becomes `noinline` (emitted alongside the convergent group and any `!dbg` scope), so `opt` and libNVVM honor the request. An `#[inline(never)]` helper is a real PTX function call. The one exception is the `#[inline(never)]` that `#[device]` adds to a helper with its own type or const parameters: it only keeps each copy visible to the device collector on the host side and is not carried to device code, so const-generic arguments keep folding into their kernels. | | Standalone `#[device]` Functions | **Full** | Device functions compiled without any kernel present. Clean export names for C++ consumption. | | Multi-Kernel Modules | **Full** | Multiple `#[kernel]` functions in a single source file compile to a single PTX module. | @@ -181,7 +199,7 @@ Anonymous promoted allocations remain unsupported. | Feature | Status | Description | |:--------|:-------|:------------| -| `ptx_asm!` Macro | **Partial** | CUDA inline PTX with `%0` operands, `in`, `out`, and `inout`; up to 16 output operands across `out` with `=`-prefixed constraints and `inout` with `+`-prefixed constraints; up to 16 explicit inputs; CUDA register constraints `h`, `r`, `l`, `q`, `f`, and `d`; immediate integer constraint `n`; compile-time string constraint `C`; `clobber("memory")`; and `options(register_only)` for pure register snippets. With multiple output operands, the macro writes tuple results back in declaration order. By default, snippets are treated as side-effecting and stay inside their current control flow. Use `options(register_only, may_diverge)` only for pure snippets that are safe to move across divergent control flow; **never** use it for `.sync` instructions or collectives. More than 16 output operands are not implemented yet. | +| `ptx_asm!` Macro | **Partial** | CUDA inline PTX with `%0` operands, `in`, `out`, and `inout`; up to 16 output operands across `out` with `=`-prefixed constraints and `inout` with `+`-prefixed constraints; up to 16 explicit inputs; CUDA register constraints `h`, `r`, `l`, `q`, `f`, and `d`; immediate integer constraint `n`; compile-time string constraint `C`; `clobber("memory")`; and `options(register_only)` for pure register snippets. With multiple output operands, the macro writes tuple results back in declaration order. By default, snippets are treated as side-effecting and stay inside their current control flow. Use `options(register_only, may_diverge)` only for pure snippets that are safe to move across divergent control flow; **never** use it for `.sync` instructions or collectives. More than 16 output operands are not implemented yet. For how each option orders against `sync_threads()`, see {ref}`Compiler ordering around the barrier `. | --- @@ -228,7 +246,7 @@ Anonymous promoted allocations remain unsupported. |:--------|:-------|:------------| | Warp Shuffle Operations | **Full** | `shuffle`, `shuffle_xor`, `shuffle_down`, `shuffle_up`. Unsuffixed forms take `u32`; `_f32`, `_u64`, `_f64` variants and a `_sync` form of each. | | Warp Vote Operations | **Full** | `all(pred)`, `any(pred)`, `ballot(pred)` → bitmask. | -| Lane/Warp ID | **Full** | `lane_id()` (0–31), `warp_id()`. Direct register reads. | +| Lane/Warp ID | **Full** | `lane_id()` (0–31) reads `%laneid`; `warp_id()` is `threadIdx.x / 32`, valid for 1D blocks only. | | Warp Reduction (`redux.sync`) | **Full** | One-instruction full-warp reduction. Integers on sm_80+: `redux_sync_add`, `_min_{u32,i32}`, `_max_{u32,i32}`, `_and`, `_or`, `_xor`. `f32` min/max with optional `.abs` and `.NaN` on `sm_100a`/`sm_100f`/`sm_103a`/`sm_103f`. No `f64` form. | ## Runtime Library: Cooperative Groups @@ -293,6 +311,6 @@ Anonymous promoted allocations remain unsupported. | Dynamic Dispatch (`dyn Trait`) | **N/A** | Use generics with static dispatch. Haven't found a real need for this. | | Heap Allocation (`Box`, `Vec`) | **N/A** | CUDA has a device-side heap (`malloc`/`free` in kernels), and the compiler allows the `alloc` crate through -- but no device-side `#[global_allocator]` is wired up today. Even if it were, device `malloc` is extremely slow (serialized, fragmented, uncoalesced). Use slices and `SharedArray`. | | `String` / `format_args!` | **N/A** | Use `gpu_printf!` for formatted output. | -| Panic / Unwinding | **N/A** | Panic paths exist in MIR but the compiler strips `core::panicking::*` and all unwind edges. The GPU hardware *can* support unwinding (absolute branches + per-thread call stack tracking post-Volta), but the CUDA toolchain (nvcc/ptxas) doesn't expose it today -- no landing pads survive to PTX. If a panic path is reached at runtime the GPU traps (same as `panic=abort`). NVIDIA has an active project to add C++ exception support to CUDA for automotive safety; the current cuda-oxide design is forward-compatible with that work. Use `gpu_assert!()` + `trap()` for explicit runtime checks today. | +| Panic / Unwinding | **N/A** | Panic paths exist in MIR but the compiler strips `core::panicking::*` and all unwind edges. The GPU hardware *can* support unwinding (absolute branches + per-thread call stack tracking post-Volta), but the CUDA toolchain (nvcc/ptxas) doesn't expose it today -- no landing pads survive to PTX. If a panic path is reached at runtime the GPU traps (same as `panic=abort`). NVIDIA has an active project to add C++ exception support to CUDA for automotive safety; the current cuda-oxide design is forward-compatible with that work. Use `gpu_assert!()` + `trap()` for explicit runtime checks today. A device function declared `-> !` with a real body, such as an error store followed by `loop {}`, is called normally and the call is followed by `unreachable`; only panic entry points, `-> !` bodies that are nothing but a panic, and calls through function pointers lower to the trap. | | Standard Library (`std`/`alloc`) | **N/A** | `std` is forbidden. `alloc` is allowed by the collector but has no backing allocator. Only `core` is fully functional. `Option`, `Result`, iterators all work. | | Texture Memory | **N/A** | Lower priority given TMA availability on Hopper+. | diff --git a/cuda-oxide/cuda-oxide-book/async-programming/overlapping-transfers-and-compute.md b/cuda-oxide/cuda-oxide-book/async-programming/overlapping-transfers-and-compute.md index c8e17ecc32..820f5b14c0 100644 --- a/cuda-oxide/cuda-oxide-book/async-programming/overlapping-transfers-and-compute.md +++ b/cuda-oxide/cuda-oxide-book/async-programming/overlapping-transfers-and-compute.md @@ -72,7 +72,7 @@ completions[slot] = Some(streams[slot].record_event(None)?); ``` The complete setup, slot-reuse wait, and final drain are in the -[`run_overlapped` implementation](https://github.com/NVlabs/cuda-oxide/blob/main/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/src/main.rs). +[`run_overlapped` implementation](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap/src/main.rs). ## Measured impact @@ -92,7 +92,7 @@ On the reference RTX 4090 (`sm_89`) run: | Overlapped pinned | 99.332 ms | | Speedup | 1.76x | -The full transfer table is in the [example README](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap). +The full transfer table is in the [example README](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap). Results vary with the GPU, PCIe topology, driver, and system load. ## Practical tradeoffs @@ -103,4 +103,4 @@ before increasing the number of streams. For the lower-level stream model, see [Scheduling and Streams](scheduling-and-streams.md). For the complete runnable -implementation, see the [`pinned_overlap` example](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap). +implementation, see the [`pinned_overlap` example](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/pinned_overlap). diff --git a/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md b/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md index 60ca175d69..64dc2797df 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md +++ b/cuda-oxide/cuda-oxide-book/compiler/architecture-overview.md @@ -129,13 +129,18 @@ 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` | @@ -146,12 +151,14 @@ cuda-oxide is split into focused crates. Here is every one and its role: | `cuda-intrinsics-gen`| Extractor and deterministic source generator for the intrinsics | | `cuda-artifact-finalizer` | Driver-independent NVVM IR, LTOIR, and PTX finalization | | `oxide-artifacts` | Architecture-neutral embedded device artifact metadata | +| `oxide-tile-export` | Export descriptions shared by `cuda-macros`, the backend and `cargo-oxide`; part of the Tile IR interop work, not released yet | +| `oxide-tilebc` | Tile IR bytecode writer for those export descriptions | | `dialect-iket` | pliron dialect modelling in-kernel event tracing | | `iket-lower` | `dialect-iket` profiles + instrumentation lowering | | `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/compiler-optimizations.md b/cuda-oxide/cuda-oxide-book/compiler/compiler-optimizations.md index bcfc8fd0a1..3a2d6515af 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/compiler-optimizations.md +++ b/cuda-oxide/cuda-oxide-book/compiler/compiler-optimizations.md @@ -272,18 +272,19 @@ unroll; the loop continues with the same behavior. ## Supported loops and safe fallback -The current pass intentionally recognizes a conservative set of loop shapes: - -Both forms need a counter the pass can follow: a compile-time-constant starting -value, the same fixed step on every path back to the header, and a simple, -side-effect-free test that directly compares the counter with its bound. The -loop must also have one ordinary entry before its header. For unusual or -irreducible control flow, the compiler warns and does not unroll the loop. +The current pass intentionally recognizes a conservative set of loop shapes. +A counted `while` loop and a range `for` loop (`for i in a..b` and +`(a..b).rev()`, over any integer counter including `isize`) need a counter the +pass can follow: a compile-time-constant starting value, the same fixed step +on every path back to the header, and a simple, side-effect-free test that +directly compares the counter with its bound. The loop must also have one +ordinary entry before its header. For unusual or irreducible control flow, the +compiler warns and does not unroll the loop. | Shape | Current behavior | |:------|:-----------------| -| Explicit counted `while` loop | Supported when the requirements above hold | -| Range `for` loop (`for i in a..b`) | Supported when the requirements above hold | +| Explicit counted `while` loop | Supported, over any integer counter including `isize`, when the requirements above hold | +| Range `for` loop (`for i in a..b` and `(a..b).rev()`) | Supported, over any integer counter including `isize`, when the requirements above hold | | Other iterator `for` loops (`a..=b`, `.step_by(k)`, `.iter()`) | Not yet recognized; warns | | Constant trip count with `#[unroll]` | Fully unrolled | | Runtime, loop-invariant limit with `#[unroll(N)]` | Partially unrolled | @@ -302,7 +303,10 @@ first. To bound compile time and memory, one annotation may create at most 1,024 body copies, 8,192 cloned basic blocks, and 65,536 cloned operations. A request over -one of those limits warns and is not unrolled. +one of those limits warns and is not unrolled. Every warning about a loop that +was not unrolled names the `#[unroll]` attribute it refers to as +`file:line:column` on a `-->` line, so a kernel with several annotated loops +shows which one was left alone. For the complete kernel-author reference, see {ref}`Loop unrolling `. 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 new file mode 100644 index 0000000000..4d7976bf36 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/compiler/images/drop-proof-termination.svg @@ -0,0 +1,155 @@ + + + + + + + + + +A destructor that never returns is not a no-op +cuda-oxide erases a destructor call when a proof says it does nothing visible. The proof now also demands that the destructor returns. +RUST SOURCE (examples/drop_glue, a compile-only probe) +MAIN (the bug: the call was erased) +PR #1361 (the fix: the call is kept) + +examples/drop_glue/src/main.rs +struct LoopingDrop(u8); + +impl Drop for LoopingDrop { + #[inline(never)] + #[allow(clippy::empty_loop)] + fn drop(&mut self) { + loop {} + } +} + +#[kernel] +pub fn diverging_cfg_drop_probe(mut out: DisjointSlice<u32>) { + if thread::index_1d().get() != 0 { + return; + } + + unsafe { + let slot = out.as_mut_ptr(); + slot.write(0x1111_1111); +store A + + { + let _guard = LoopingDrop(0); +drop: never returns + } + + // Semantically unreachable. + slot.write(0x2222_2222); +store B: must never run + } +} + +What the old proof checked +It walks the destructor's blocks and looks for +visible effects: writes through pointers, asserts, +inline asm. loop {} has none, so the verdict was +harmless: erase the call, lower the drop to a branch. +It never asked whether the destructor returns. +A block seen twice was skipped (a loop looked like +a join), and a call back into a destructor already +under proof was accepted as harmless. + +verdict: no-op. The kernel IR has 0 destructor calls. + +drop_glue.ptx (main backend), thread 0 +$L__BB4_1: // %bb3 + st.global.b32 [%rd1], 572662306; +$L__BB4_2: // %bb5 + ret; +572662306 is 0x22222222: store B ran. +Store A is gone too: once the call between them was +erased, B overwrote A at once and LLVM deleted A. +No loop. The thread returns. + +What the new proof checks +The same walk, depth first, with three states +per block: Unseen, InProgress, Proven. +A successor that is still InProgress is a +cycle: the proof fails and the call is kept. +A Proven successor is a join: skip it, go on. +A call into a destructor that is already +under proof fails the same way. +loop {} fails: bb1 branches back to bb1. + +not proven: the kernel IR keeps the call. + +drop_glue.ptx (PR backend), thread 0 +$L__BB4_1: // %bb3 + st.global.b32 [%rd1], 286331153; +$L__BB4_2: // %bb1.i.i + bra.uni $L__BB4_2; +286331153 is 0x11111111: store A ran. +Then bra.uni to its own label: the thread +spins in the destructor's loop {}, inlined. +Store B is gone, as in the Rust source. + + + +The three states, on two destructors from the example +Depth-first walk from the entry block. A block turns Proven only after every successor is Proven. + +Unseen +not reached yet + +InProgress +on the current path; its exits are not yet proven + +Proven +every exit returns and does nothing visible +Fails: an InProgress block met again, or a call into +a destructor already under proof. Verdict: keep the call. +LoopingDrop::drop +bb0: goto bb1; bb1: goto bb1; + +bb0 + +bb1 + + +goto bb1 again: +still InProgress, a cycle +bb1 is met again before it is Proven. The proof fails, +the destructor call stays, the PTX spins after store A. +RecursiveDrop: the callee is already under proof, same +verdict (858993459 = 0x33333333, then a self-loop). +DiamondDrop::drop +an if/else that joins back; every exit returns + +test + +left_noop + +right_noop + +return + + + + +second visit: +Proven, a join +if self.0 == 0 { left_noop() } else { right_noop() } +return is reached twice; the second visit finds it Proven: a join. +Every path returns and writes nothing, so the drop is still erased. +PTX: two stores (0x55555555, 0x66666666) and ret, on both backends. + 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 new file mode 100644 index 0000000000..d335217568 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/compiler/images/enum-alignment-anchor.svg @@ -0,0 +1,145 @@ + + + + + +Enum alignment anchors: the zero-length anchor is capped at 16 bytes ([0 x i128]) +mir-lower raises an enum storage type's natural alignment with a zero-length integer array. Its width was align * 8 with no bound; the cap is the one unions got in #610. +THE ENUM AND ITS ANCHOR (what mir-lower mints) +BEFORE (the width grows with the alignment, then breaks) +AFTER #1357 (one cap, stamped alignment, one helper) + +#[repr(align(536870912))] +pub enum HugeAlign { A, B } +rustc: size 2^29, align 2^29, u8 tag at offset 0 + +build_enum_slot_map mints the LLVM storage type + +{ [0 x iN], i8, [size - 1 x i8] } +[0 x iN] +0 bytes: only sets the struct's natural alignment + + + +i8 +[size - 1 x i8] +tag +filler to rustc's size + + + +size = align bytes; arrays stride by it; the anchor adds none +Here the tag and filler bytes give the struct natural +alignment 1. When rustc wants more, slot 0 gets an +anchor whose element width is align * 8. It adds no +bytes, so size and array stride stay rustc's; only +LLVM's idea of the type's alignment moves. + +Unions always carry one (union_storage.rs); an enum +gets one when its byte claims under-align it +(enum_layout.rs). Both are minted once per type. + +enum_layout.rs on main: width = align * 8, no bound +let anchor_int = IntegerType::get( + ctx, (required_align * 8) as u32, Signedness::Signless); +let anchor = ArrayType::get(ctx, anchor_int.into(), 0); +llvm_fields.insert(0, anchor); +The same three lines, with align.min(16), had capped unions since #610. +What the anchor became +align(N) +anchor on main +LLVM +anchor now + +8 +[0 x i64] +ok +[0 x i64] +16 +[0 x i128] +ok +[0 x i128] +32 +[0 x i256] +16 vs 32 +[0 x i128] +64 +[0 x i512] +16 vs 32 +[0 x i128] +128 +[0 x i1024] +16 vs 32 +[0 x i128] +2^20 +[0 x i8388608] +16 vs 32 +[0 x i128] +2^21 +[0 x i16777216] +rejected +[0 x i128] +2^28 +[0 x i2147483648] +rejected +[0 x i128] +2^29 +[0 x i0] +rejected +[0 x i128] + +ok + valid; the same anchor after the fix. Align 8 and 16 are the controls. +16 vs 32 + valid IR but no alignment gained: our data layout strings stop at + i128:128, so opt saw 16; llc's default NVPTX layout has i256:256: 32. +rejected + error: bitwidth for integer or byte type out of range (opt and llc), + build exit 101. From 2^21 up the width passes LLVM's 2^23-bit limit + (i8388608 is the last accepted); at 2^29, (2^29 * 8) as u32 == 0. + +the cap: anchor_align = min(align, 16) +main { [0 x i0], i8, [536870911 x i8] } +PR { [0 x i128], i8, [536870911 x i8] } +size 2^29, stride 2^29, tag at offset 0: unchanged. +Builds, runs and passes memcheck; main exits 101. + +the Rust alignment rides on the ops, not the anchor +alloca [3 x { [0 x i128], i8, [63 x i8] }], align 64 +store { [0 x i128], i8, [63 x i8] } .., align 64 +load i8, ptr %v13, align 536870912 +.local .align 64 .b8 __local_depot1[192] +.param .align 64 .b8 by_value_tag64_param_0[64] +shl.b64 %rd18, %rd3, 29 stride 2^29 +Stamped on every alloca, load, store, spill, kernel +parameter and return typed as the enum. PTX for align +32, 64 and 128 is instruction-identical to main. + +one helper for unions and enums +pub(crate) const MAX_NATURAL_ALIGN: u64 = 16; +pub(crate) fn alignment_anchor_type( + ctx: &mut Context, abi_align: u64) -> TypeHandle +Called from union_storage.rs and enum_layout.rs. +16 is where the three data layout strings we emit, +llvm_type_size_align and the exporter's +natural_alignment all stop (i128). + + + +What changes and what does not +Changes: only the anchor's element width for alignments of 32 and up (i256, i512, i1024, i8388608 and i0 all become i128). Unchanged: size, stride, and the alignment +stamped on the memory operations. Enums aligned 2^20, 2^21 and 2^29 go from exit 101 on main to building and running; align 32 to 128 keep the same PTX. A wider anchor never bought alignment. + 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 93217b0277..668013f72a 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/lowering-pipeline.md +++ b/cuda-oxide/cuda-oxide-book/compiler/lowering-pipeline.md @@ -229,6 +229,21 @@ It is not elegant, but it is exactly how C compilers have handled tagged unions for decades. LLVM's optimizer is quite good at cleaning up the redundant insertvalue/extractvalue chains. +An enum whose tag and payload bytes under-align its storage, a `repr(align(N))` +enum for example, also gets a zero-length integer array at slot 0 as an alignment +anchor. The anchor adds no bytes and is capped at 16 bytes; the Rust alignment +itself is stamped on every alloca, load and store of the enum. + +```{figure} images/enum-alignment-anchor.svg +:align: center +:width: 100% + +The enum alignment anchor. A zero-length integer array at slot 0 raises the storage type's natural +alignment without adding bytes. Its width used to be `align * 8` with no bound, which LLVM rejects +from align 2^21 up and which wraps to `i0` at 2^29. It is now capped at `[0 x i128]`, as union +storage has been since #610, while the Rust alignment rides on every memory operation. +``` + --- ## GPU Intrinsic Conversion @@ -572,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-08-28`) 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/mir-importer.md b/cuda-oxide/cuda-oxide-book/compiler/mir-importer.md index 86fb1d2e7b..d6c3281ade 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/mir-importer.md +++ b/cuda-oxide/cuda-oxide-book/compiler/mir-importer.md @@ -376,6 +376,20 @@ slip through (array bounds checks, unwrap on `None`), a GPU trap is the correct behavior anyway -- there is nothing useful a GPU thread can do to "recover" from a logic error mid-kernel. +Drop glue follows the same rule. The importer erases a destructor call only when +its proof shows that the destructor does nothing visible and returns. A destructor +that loops forever, or calls back into an instance already under proof, keeps its +call, so the code after it never runs, as in the Rust source. + +```{figure} images/drop-proof-termination.svg +:align: center +:width: 100% + +The drop no-op proof requires that the destructor returns. Before, the call to a `loop {}` +destructor was erased and the store after it ran. Now the call stays, so the PTX stores once and +spins, while an `if`/`else` destructor that joins back is still erased. +``` + --- ## Putting It All Together diff --git a/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md b/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md index c3057e5280..6225db2323 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md +++ b/cuda-oxide/cuda-oxide-book/compiler/mlir-dialects.md @@ -195,7 +195,7 @@ handles all of that flattening. ### Operations -At the pinned `pliron` revision the dialect defines 69 operations: +The pinned `pliron-llvm` 0.20.0 release defines 69 operations: | Category | Examples | Count | | :----------- | :-------------------------------------------------------------------------------------------------------------------------------------- | ----: | @@ -216,9 +216,9 @@ At the pinned `pliron` revision the dialect defines 69 operations: `llvm-export` adds one operation of its own on top of those, `llvm.dbg_value`, alongside the address-space and fp16 helpers mentioned above. -Because the dialect is upstream, this table moves when the `pliron` pin moves -rather than when this repository changes; `pliron-llvm`'s `src/ops.rs` is the -list it is counting. +The operation inventory follows the pinned upstream release. This table counts +the operations in [`pliron-llvm`'s `src/ops.rs`](https://github.com/pliron-org/pliron/blob/435536c8611af34c48c8f73d78e77894a05d103d/pliron-llvm/src/ops.rs), +the source revision recorded in the 0.20.0 crate. If you have read LLVM IR before, nothing here will surprise you. The operation names are intentionally the same as their LLVM counterparts, prefixed with @@ -303,16 +303,16 @@ they become `call` instructions to `@llvm.nvvm.*` intrinsics. ### Architecture Coverage -At catalog SHA-256 `016c9d9b` (the stamp in every `ops/generated/` file -header), the dialect holds 576 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 `016c9d9b`, 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,22 +321,22 @@ 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 550 operations, resolved from 1029 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 | | :-------------------------- | :---------------------------------------------------------------------------- | --: | | Tensor Core Gen 5 + TMEM | `tcgen05` | 210 | -| Tensor Memory Accelerator | `tma` | 111 | +| Tensor Memory Accelerator | `tma` | 116 | | Special registers | `sreg` | 44 | -| Packed (SIMD-in-register) | `packed_alu`, `packed_conversion`, `packed_atomic` | 51 | -| Async copy and barriers | `cp_async`, `mbarrier_extended`, `mbarrier_basic`, `sync` | 35 | +| Packed (SIMD-in-register) | `packed_alu`, `packed_conversion`, `packed_atomic` | 55 | +| 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/pliron.md b/cuda-oxide/cuda-oxide-book/compiler/pliron.md index 99f1c38751..ec8d2b1bd6 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/pliron.md +++ b/cuda-oxide/cuda-oxide-book/compiler/pliron.md @@ -73,7 +73,7 @@ model -- operations, regions, basic blocks, types, attributes, dialects, passes | DSLs required | None -- just Rust macros | TableGen, ODS | | Debugging | Standard Rust tooling (`dbg!`, rust-gdb) | gdb on C++ templates | | Extensibility | Add a dialect as a Rust crate | C++ extension against LLVM headers | -| Dependency weight | One crate (git dependency) | Gigabytes of LLVM build artifacts | +| Dependency weight | One published crate | Gigabytes of LLVM build artifacts | For cuda-oxide, this means the entire compiler -- from MIR import to LLVM IR export -- stays within the same Rust build and debugging environment as the rest diff --git a/cuda-oxide/cuda-oxide-book/compiler/rustc-codegen-cuda.md b/cuda-oxide/cuda-oxide-book/compiler/rustc-codegen-cuda.md index b9513030e5..683605591e 100644 --- a/cuda-oxide/cuda-oxide-book/compiler/rustc-codegen-cuda.md +++ b/cuda-oxide/cuda-oxide-book/compiler/rustc-codegen-cuda.md @@ -215,7 +215,7 @@ takes the `#[kernel]` macro's base name and, for generic and closure- generic instantiations, appends `_TID_` where `` is rustc's 128-bit type-id hash of the concrete generated kernel function item. That `FnDef` contains ordered type and const arguments. The host launcher -computes the same hash through the `core::intrinsics::type_id` intrinsic +computes the same hash through the `core::intrinsics::reflection::type_id` intrinsic (wrapped by `cuda_host::type_id_u128_of_val`), so both sides agree byte-for-byte within one unified build. Non-generic kernels keep their bare base name. diff --git a/cuda-oxide/cuda-oxide-book/compiler/rustc-public.md b/cuda-oxide/cuda-oxide-book/compiler/rustc-public.md index 737cb0fdb5..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-08-28" +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/conf.py b/cuda-oxide/cuda-oxide-book/conf.py index e8d2f3f5ff..a69dc4435d 100644 --- a/cuda-oxide/cuda-oxide-book/conf.py +++ b/cuda-oxide/cuda-oxide-book/conf.py @@ -50,7 +50,7 @@ "icon_links": [ { "name": "GitHub", - "url": "https://github.com/NVlabs/cuda-oxide", + "url": "https://github.com/NVIDIA/cuda-rust", "icon": "fa-brands fa-github", "type": "fontawesome", }, 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 6c0325b9bc..cec0a77fba 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-08-28 install --git https://github.com/NVlabs/cuda-oxide.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 ``` 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 new file mode 100644 index 0000000000..27a7d1e396 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/getting-started/images/backend-install-by-rename.svg @@ -0,0 +1,108 @@ + + + + + +Installing the cached backend by rename, not in place +Every rustc that compiles a kernel maps the shared 428 MB librustc_codegen_cuda.so. std::fs::copy rewrote that inode under them; a sibling temporary plus rename leaves it alone. +THE SETUP (one file, many readers, three writers) +BEFORE (fs::copy rewrites the mapped inode) +AFTER (copy beside it, rename over the name) + + + +Who installs the backend +cargo oxide setup publish_to_cache +a standalone project's rebuild build_and_cache +clone of main, no dependency auto_fetch_and_build +all three call install_backend_into, at any time + +install_backend_into(cache_dir, built_so, ..) + +~/.cargo/cuda-oxide/ +librustc_codegen_cuda.so 428 MB, inode A +source-rev.txt the commit it came from +toolchain-fingerprint.txt the rustc that built it + +dlopen at start, mapped MAP_PRIVATE until exit + +every rustc that compiles a kernel +rustc -Zcodegen-backend=<the cached .so> ... +holds the mapping for the whole compile; ld.so's +relocations live in private pages of that inode +A reinstall can land while such a rustc runs: +setup in one shell, a kernel build in another. + +std::fs::copy(built_so, &so_path) +opens the destination with truncate(true): +inode A is cut to 0 bytes, then rewritten + + +what rustc's mapping of inode A sees +rustc's mapping: 428 MB, MAP_PRIVATE, until it exits + + + + +inode A before: all 428 MB present + +new bytes + +gone: a touch here is SIGBUS +inode A is cut to 0 bytes, then refilled by the copy; +the mapping's private pages (relocations) go too: +SIGSEGV even when the new bytes are identical + + +the running rustc dies +SIGBUS when it touches a page past the new end +SIGSEGV through relocations the truncate threw +away, even when the new bytes are identical +A rustc starting during the 0.1 to 3 s copy +opens a half-written file. + +install_backend_into, all three writers +let tmp = temporary_backend_path(cache_dir); +// = librustc_codegen_cuda.so.<pid>.<nonce>.tmp +fs::copy(built_so, &tmp)?; // old .so and record stay +record_source_rev(cache_dir, None); // follow-up 1 +fs::rename(&tmp, &so_path)?; // atomic, same directory +write_toolchain_fingerprint(..); record_source_rev(..); +// any error: remove tmp, return it; the old .so stays + + +two inodes, one name +rustc's mapping stays on inode A until it exits + + + + +inode A: unnamed now, alive until its last mapper exits + +inode B: named librustc_codegen_cuda.so, already complete +a rustc starting now opens inode B, complete + +Two follow-ups +1 the record is cleared just before the rename, so + a failed copy keeps the old backend and its record +2 cargo oxide update sweeps a .tmp left by an install + killed between copy and rename; installs never do + +Why the rename is safe and the copy was not +rename(2) swaps the directory entry. A mapped inode stays allocated until its last mapper exits, and a new open sees the complete new file. +std::fs::copy reopened that inode with O_TRUNC under every mapping: pages past the new end are gone (SIGBUS), and so are the mapping's private pages (SIGSEGV). +Measured: a rustc with the 428 MB backend mapped, killed 5/5 in place (identical bytes), finished 5/5 by rename; an mmap reader of a 64 MB install, SIGBUS 6/6 before, 0/6 after. + diff --git a/cuda-oxide/cuda-oxide-book/getting-started/installation.md b/cuda-oxide/cuda-oxide-book/getting-started/installation.md index 59ae7b28ae..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-08-28 -rustup component add rust-src rustc-dev rust-analyzer clippy rustfmt llvm-tools --toolchain nightly-2026-08-28 +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,14 +306,23 @@ 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-08-28 install --git https://github.com/NVlabs/cuda-oxide.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. +```{figure} images/backend-install-by-rename.svg +:align: center +:width: 100% + +Installing the cached backend by rename. Every `rustc` that compiles a kernel maps the one shared `.so`. +The installer copies a new build to a sibling temporary and renames it over the name, so running +compilers keep the old inode and a compiler that starts during the install opens a complete file. +``` + To discover the examples available in a cuda-oxide checkout: ```bash @@ -337,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: @@ -396,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-08-28`, 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-08-28`. +- `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/error-handling-and-debugging.md b/cuda-oxide/cuda-oxide-book/gpu-programming/error-handling-and-debugging.md index 61c638b884..f48e28e4b5 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 @@ -51,6 +51,10 @@ pub fn debug_kernel(data: &[f32], mut out: DisjointSlice) { } ``` +A placeholder without a type character prints its argument by type: the +`usize` index above goes out as `%llu` and the `f32` element as `%f`, while +`{:x}` or `{:.2f}` cast the argument to the C type they name. + ### Important details - **Flush requires sync.** Output is buffered on the GPU and only appears on @@ -484,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/images/diverging-call-lowering.svg b/cuda-oxide/cuda-oxide-book/gpu-programming/images/diverging-call-lowering.svg new file mode 100644 index 0000000000..9a0932eb90 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/images/diverging-call-lowering.svg @@ -0,0 +1,112 @@ + + + + + +A device function that never returns is now called, not replaced by a trap +Such a function may store an error code and then loop. The importer replaced the call with trap; and lost the store. Now it emits the call, then unreachable, as rustc's own backend does. +RUST SOURCE (the function_item_call example) +BEFORE, ON MAIN (the call became a trap) +AFTER, WITH #1364 (the call is kept) + +write_then_diverge (src/main.rs) +#[device] +unsafe fn write_then_diverge(ptr: *mut u32) -> ! { + unsafe { + ptr.write(0x1111_1111); + } + loop {} +} +Returns `!`: it never comes back. The store is the +work that has to survive. + +in the #[kernel] diverging_callable +if flag == 3 { + unsafe { + write_then_diverge(out.as_mut_ptr()); + } +} +An ordinary direct call. Nothing follows it: the +callee's type promises it does not return. +Both backends get the same MIR from rustc: a Call with +no return edge (target: None), because the callee +never returns. The middle and right columns differ +only in what the importer emits for that call. + +What the programmer expects +out[0] = 0x1111_1111, then the thread spins forever. +An ordinary call, as on the host, and as nvcc treats +a __noreturn__ device function. + +LLVM IR before opt (@diverging_callable) +bb6: + %v14 = extractvalue { ptr, i64 } %v6, 0 + call void asm sideeffect "trap;", ""() + unreachable +The call to write_then_diverge is gone. bb2 and bb4, +the other two diverging call sites, look the same. + +opt -O2, then llc -march=nvptx64 + +PTX (.entry diverging_callable) +$L__BB1_5: // %bb6 + // begin inline asm + trap; + // end inline asm +trap; appears three times in this kernel, once per +diverging call site. .func write_then_diverge is in +the module, marked .noreturn, and nothing calls it. + +On the GPU, flag == 3: no store. The thread hits trap; +and the launch fails. The value never lands in memory. + +LLVM IR before opt (same block) +bb6: + %v14 = extractvalue { ptr, i64 } %v6, 0 + call void @write_then_diverge(ptr %v14) #0 + unreachable +The call stays. unreachable only records that control +never comes back to this block. + +opt -O2, then llc -march=nvptx64 + +PTX (.entry diverging_callable) +$L__BB0_7: // %bb6 + st.global.b32 [%rd1], 286331153; +$L__BB0_8: // %bb1.i + bra.uni $L__BB0_8; +opt inlined the callee: the store (286331153 is +0x1111_1111), then the self-loop. No trap; on this path. +.func write_then_diverge is kept too, marked .noreturn. + +On the GPU, flag == 3: out[0] = 0x1111_1111, then the +thread spins. Seen on an RTX 5090 with a scratch kernel. + +Still lowered to trap; (by design) +1 panic entry points in core::panicking: + panic!(..), x.unwrap() on None, a[i] out of range +2 a -> ! body that is only a panic; not collected: + fn panic_only() -> ! { panic!("not collected") } +3 a call through a function pointer, not resolved: + let f: unsafe fn(*mut u32) -> ! = helper; f(p) + + + +How the importer decides (translate_call, for a MIR Call with no return edge) +run_pipeline records the mangled name of every function it emits for this build. translate_call resolves the callee and checks that set: +in the set: emit the call, then unreachable, which is what rustc's own LLVM backend emits for a call with no destination. +not in the set: trap; then unreachable, as before. That blanket trap was written for panic entry points, which have no body on the GPU; it swallowed user functions as well. + 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 new file mode 100644 index 0000000000..cd048c5329 --- /dev/null +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/images/inline-intent-mapping.svg @@ -0,0 +1,125 @@ + + + + + +An inline request now reaches the GPU as the matching LLVM attribute +Only #[inline(always)] used to survive. Now #[inline] is inlinehint, #[inline(always)] is alwaysinline and #[inline(never)] is noinline, as in rustc's own backend, so a never helper stays a real PTX function. +RUST SOURCE (the inline_intent example) +RECORDED, THEN LOWERED (mir.func, then the define line) +WHAT OPT AND LLC MAKE OF IT (PTX, sm_120a) + +one helper of each kind +#[inline(never)] +pub fn never_add(a: f32, b: f32) -> f32 { + a + b +} +#[inline] +pub fn hint_mul(a: f32, b: f32) -> f32 { + a * b +} +#[inline(always)] +pub fn always_sub(a: f32, b: f32) -> f32 { + a - b +} +src/main.rs: plain functions the collector +reaches from the kernel. + +inside the #[cuda_module] +#[inline(never)] +#[device] +pub fn never_dev(a: f32, b: f32) -> f32 { + a * 2.0 + b +} +#[device] +pub fn scale_by<const N: u32>(x: f32) -> f32 { + x * (N as f32) +} +scale_by has its own const parameter, so +#[device] adds an #[inline(never)] of its own +there, for the host-side collector. + +the kernel's arithmetic +let mut v = never_add(x, 1.0); +v = hint_mul(v, 2.0); +v = always_sub(v, 1.0); +v += kernels::never_dev(x, 1.0); +v + kernels::scale_by::<4>(x) +Reached through one #[inline(always)] helper. +MIR inlining is off for this package, so every +helper reaches LLVM as its own definition. + +the request, as the pipeline carries it +source mir.func LLVM +#[inline] inline_intent=hint inlinehint +#[inline(always)] inline_intent=always alwaysinline +#[inline(never)] inline_intent=never noinline +no attribute none none +rustc-codegen-cuda reads codegen_fn_attrs and sets a typed +InlineIntentAttr on mir.func. mir-lower names the keyword in +one place. A kernel entry never carries one: verifier error. + +LLVM IR before opt (inline_intent.ll) +define float @inline_intent__never_add(..) noinline #0 +define float @inline_intent__hint_mul(..) inlinehint #0 +define float @inline_intent__always_sub(..) alwaysinline #0 +define float @never_dev(..) noinline #0 +define float @scale_byKm4_EB4_(float %v0) #0 +(..) stands for (float %v0, float %v1). The last line is +scale_by::<4>: no keyword, on purpose (next box). + +The one request that is not carried +#[device] puts #[inline(never)] on a helper with its own +type or const parameters so each monomorphization stays +visible to the host-side collector. That is a host need, +not a device intent, so the backend drops it and N keeps +folding into the kernel. A user's own attribute on such a +helper still does not reach device code: #1471 stays open. + +PTX, inside .entry intent (opt -O2, then llc) +call.uni (retval0), inline_intent__never_add, (param0); +call.uni (retval0), never_dev, (param0, param1); +fma.rn.f32 %r13, %r8, 0f40800000, %r12; +Two real calls, one per noinline helper. 0f40800000 is +4.0: scale_by::<4> folded into the entry. always_sub +appears nowhere, alwaysinline left no function. hint_mul +was inlined by opt, as the hint allows. + +.func headers left in the module +.func (..) inline_intent__never_add( +.visible .func (..) never_dev( +.visible .func (..) scale_byKm4_EB4_( +(..) stands for (.param .b32 func_retval0). scale_by stays +only because every #[device] function is kept alive through +@llvm.used; the entry does not call it. + +RTX 5090: SUCCESS, 256 of 256 match the host reference. +verify-code-shape.sh checks these lines in compile-only CI. + +Before, on main +No keyword on any of these define lines. opt inlined +the never helpers like the rest, so a helper that asked +to stay a separate GPU function never did, and #[inline] +was dropped. Only alwaysinline came through, as a string +attribute the exporter matched by name. + + + +The libNVVM route needed one more piece (nvvm-transforms, legalize_pointer_returns) +A noinline helper that returns a shared-memory pointer is now a real function. libNVVM declares that 32-bit return slot as .param .b64 while the body stores 32 bits and the call site reads 32, +so ptxas rejected addressof_sharedarray on --emit-nvvm-ir. The pass gives such a definition a generic pointer return and an addrspacecast on both sides of each call; parameters, aggregates, +kernels and declarations are untouched, and a function whose address is taken is reported. The libNVVM defect is tracked with its team; the pass goes when the fix ships. + diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/kernel-families.md b/cuda-oxide/cuda-oxide-book/gpu-programming/kernel-families.md index c0e8cd45f8..64e4c4cb86 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/kernel-families.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/kernel-families.md @@ -992,10 +992,10 @@ compiled shape is incompatible with the matrix dimensions. See: -* [`gemm_sol_final`](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final) -* [`KernelFamily` integration tests](https://github.com/NVlabs/cuda-oxide/blob/main/cuda-oxide/crates/cuda-host/tests/kernel_family.rs) -* [`cuda-host` kernel-family implementation](https://github.com/NVlabs/cuda-oxide/blob/main/cuda-oxide/crates/cuda-host/src/kernel_family.rs) -* [`cuda-host` README](https://github.com/NVlabs/cuda-oxide/blob/main/cuda-oxide/crates/cuda-host/README.md) +* [`gemm_sol_final`](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final) +* [`KernelFamily` integration tests](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-oxide/crates/cuda-host/tests/kernel_family.rs) +* [`cuda-host` kernel-family implementation](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-oxide/crates/cuda-host/src/kernel_family.rs) +* [`cuda-host` README](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-oxide/crates/cuda-host/README.md) ## Design guidelines diff --git a/cuda-oxide/cuda-oxide-book/gpu-programming/kernels-and-device-functions.md b/cuda-oxide/cuda-oxide-book/gpu-programming/kernels-and-device-functions.md index 2580a40058..303334a50e 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/kernels-and-device-functions.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/kernels-and-device-functions.md @@ -191,6 +191,30 @@ pub fn magnitude(x: f32, y: f32) -> f32 { | Callable from device code| Yes | Yes | Yes | | Annotation required | Always | Only for standalone/cross-crate/FFI | Never | +### Inlining helpers + +A helper's `#[inline]` request reaches the GPU the same way it reaches the host +backend. `#[inline]` becomes the LLVM `inlinehint`, `#[inline(always)]` +becomes `alwaysinline`, and `#[inline(never)]` becomes `noinline`, so a helper +marked `never` stays a separate `.func` in the PTX and is reached with a real +`call.uni`. One exception: the `#[inline(never)]` that `#[device]` adds to a +helper with its own type or const parameters is a host-side collector boundary, +not a device intent, so it is not carried and const-generic arguments keep +folding into the kernel. + +```{figure} images/inline-intent-mapping.svg +:align: center +:width: 100% + +How an inline request reaches the GPU. Left, one helper of each kind from the +`inline_intent` example. Middle, the request recorded on `mir.func` as a typed +`inline_intent` attribute and lowered to the LLVM function attribute in one +place. Right, what opt and llc make of it: `call.uni` for the `noinline` +helpers, no function left for the `alwaysinline` one, and the const-generic +factor folded to an immediate. Before this mapping only `alwaysinline` +survived, and LLVM inlined a `never` helper anyway. +``` + ## What Rust works on the GPU cuda-oxide compiles standard Rust through `rustc` -- it is not a subset @@ -321,8 +345,9 @@ pub fn sum_four(mut out: DisjointSlice) { ``` The pass currently recognizes explicit counted `while` loops and range `for` -loops (`for i in a..b`). Other iterator loops, such as `a..=b` or -`.step_by(k)`, are not yet recognized. +loops (`for i in a..b` and `(a..b).rev()`) over any integer counter, including +`isize`. Other iterator loops, such as `a..=b` or `.step_by(k)`, are not yet +recognized. Use `#[unroll(N)]`, where `N >= 2`, when the trip count is only known at runtime. The loop then does `N` iterations' work per trip. A small remainder loop handles @@ -351,6 +376,9 @@ unroll that loop. Partial unrolling also requires a counted-up loop: the counter must have a positive step, use `<` or `<=`, and compare against a limit that does not change inside the loop. The compiler warns and does not unroll unsupported requests. +Every warning about a loop that was not unrolled names the `#[unroll]` +attribute it refers to as `file:line:column` on a `-->` line, so a kernel with +several annotated loops shows which one was left alone. To keep generated code bounded, one annotation may create at most 1,024 body copies, 8,192 cloned basic blocks, and 65,536 cloned operations. A larger @@ -512,6 +540,25 @@ In practice this means: once the call is dropped. Reach for `gpu_printf!` before the check when you need to see the values, or `gpu_assert!` for an explicit check. +A device function declared `-> !` with a real body is different. It never +returns to its caller, for example because it writes an error code and then +runs `loop {}`. The compiler calls it like any other function and places +`unreachable` after the call, so the work it does before diverging is kept. +Three cases still lower to the device trap instead: the panic entry points in +`core::panicking`, a `-> !` body that is nothing but a panic, and a call +through a function pointer. + +```{figure} images/diverging-call-lowering.svg +:align: center +:width: 100% + +How a call to a `-> !` device function lowers. Left, the Rust source: `write_then_diverge` stores a +value and then loops. Middle, what the importer used to emit: the call replaced by `trap;`, so the +store never happened. Right, what it emits now: the call, then `unreachable`, so the PTX stores the +value and spins. Panic entry points, panic-only `-> !` bodies and calls through function pointers +still lower to the trap. +``` + :::{seealso} The [Error Handling and Debugging](error-handling-and-debugging.md) chapter covers `gpu_printf!`, `gpu_assert!`, and `cargo oxide debug` for diagnosing 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 0bea99254b..8a6b94844f 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-programming/launching-kernels.md +++ b/cuda-oxide/cuda-oxide-book/gpu-programming/launching-kernels.md @@ -441,7 +441,7 @@ generated beside the kernel function. ### Complete example The -[`policy_config`](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config) +[`policy_config`](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/policy_config) example defines two policies for one generic kernel and verifies that they produce: 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 b314852322..d7bbc9a06f 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` in cutile-rs](https://github.com/NVlabs/cutile-rs/blob/main/cuda-core/tests/simt_vmm_p2p.rs). +[`cuda-core/tests/simt_vmm_p2p.rs`](https://github.com/NVIDIA/cuda-rust/blob/main/cuda-core/tests/simt_vmm_p2p.rs). The corresponding implementations are: -* [`cuda-core/src/simt/vmm.rs` in cutile-rs](https://github.com/NVlabs/cutile-rs/blob/main/cuda-core/src/simt/vmm.rs) -* [`cuda-core/src/simt/peer.rs` in cutile-rs](https://github.com/NVlabs/cutile-rs/blob/main/cuda-core/src/simt/peer.rs) +* [`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) ::: ## VMM and P2P solve different problems diff --git a/cuda-oxide/cuda-oxide-book/gpu-safety/bounds-checks.md b/cuda-oxide/cuda-oxide-book/gpu-safety/bounds-checks.md index 66c56c57d5..3516cf83c4 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-safety/bounds-checks.md +++ b/cuda-oxide/cuda-oxide-book/gpu-safety/bounds-checks.md @@ -118,7 +118,7 @@ load a, load b, fma, advance, compare, branch ``` -- exactly the shape of a hand-written raw-pointer loop. The -[`gemm_views`](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views) +[`gemm_views`](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_views) example proves this literally: it compiles each safe kernel next to an `unsafe` raw-pointer twin and verifies that the PTX has the same branch counts, the same memory instructions, no traps in the loop, and @@ -264,7 +264,7 @@ Scope, precisely: helper function keeps its own checks. - Turn it on only after the checked build passes, ideally under `compute-sanitizer`. The - [`unchecked_indexing`](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing) + [`unchecked_indexing`](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/unchecked_indexing) example shows the PTX with and without the flag. --- diff --git a/cuda-oxide/cuda-oxide-book/gpu-safety/the-safety-model.md b/cuda-oxide/cuda-oxide-book/gpu-safety/the-safety-model.md index 1be518fb87..ed7f98b663 100644 --- a/cuda-oxide/cuda-oxide-book/gpu-safety/the-safety-model.md +++ b/cuda-oxide/cuda-oxide-book/gpu-safety/the-safety-model.md @@ -198,7 +198,7 @@ in elements and must match the buffer's layout. The buffer length does not have to be a multiple of the row width; only complete tiles are returned. See the runnable -[`proof_carrying_views`](https://github.com/NVlabs/cuda-oxide/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views) +[`proof_carrying_views`](https://github.com/NVIDIA/cuda-rust/tree/main/cuda-oxide/crates/rustc-codegen-cuda/examples/proof_carrying_views) example for 1-D and 2-D versions. ### How `index_2d` is type-safe 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/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..5094602dd0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/cfg.rs @@ -0,0 +1,293 @@ +/* + * 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, OpId}, + 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 a scope holds any operation of one semantic family. +pub(super) fn scope_has_any(ctx: &Context, scope: &[Ptr], ids: &[OpId]) -> bool { + scope + .iter() + .any(|operation| ids.contains(&Operation::get_opid(*operation, ctx))) +} + +/// 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..fe2794e177 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/epilogue.rs @@ -0,0 +1,697 @@ +/* + * 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::{LocalCells, is_semantic_pointer_carrier_cast}, + 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, + cells: &LocalCells, + 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, cells, 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, cells: &LocalCells, 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; + } + cells.closed_cell(ctx, load.address_opd(ctx)).is_some() +} + +/// 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, + cells: &LocalCells, + operation: Ptr, +) -> bool { + let opid = Operation::get_opid(operation, ctx); + is_epilogue_writer_plumbing(ctx, cells, 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 cells = LocalCells::new(ctx, all_ops); + 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, &cells, *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, &cells, 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..f154312134 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/mma.rs @@ -0,0 +1,803 @@ +/* + * 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, + cfg::scope_has_any, + ghosts::smem_mma_semantic_ids, + invalid, + provenance::{ + CarrierFact, LocalCells, OrdinaryFacts, branch_successor_operands, + is_semantic_pointer_carrier_cast, 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), +} + +impl CarrierFact for MmaFact { + fn is_opaque(&self) -> bool { + matches!(self, Self::Scalar(_) | Self::Unknown(_)) + } +} + +/// 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 cells = LocalCells::new(ctx, operations); + 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, &cells, &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, + &cells, + &mut facts, + MmaFact::Unknown, + ); + let propagated = propagate_ordinary_plumbing(ctx, operations, &cells, &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. +#[allow(clippy::too_many_arguments)] +fn mma_ordinary_use_forwards_fact( + ctx: &Context, + cells: &LocalCells, + facts: &OrdinaryFacts, + value: Value, + path: &[u32], + 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), 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(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 path.first().copied() == Some(index.0) { + return true; + } + return mma_fact_at_path(facts, operation.get_result(0), path, fact); + } + if operand_index == 1 { + let mut destination_path = vec![index.0]; + destination_path.extend_from_slice(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 path.first().copied() != Some(index.0) { + return true; + } + return mma_fact_at_path(facts, operation.get_result(0), &path[1..], fact); + } + if opid == MirStoreOp::get_opid_static() && operation.get_num_operands() == 2 { + let store = MirStoreOp::new(user); + if operand_index == 0 { + return cells.closed_cell(ctx, store.address_opd(ctx)).is_some(); + } + if operand_index == 1 + && let Some((root, mut destination_path)) = + cells.closed_cell(ctx, store.address_opd(ctx)) + { + destination_path.extend_from_slice(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((_, cell_path)) = cells.closed_cell(ctx, load.address_opd(ctx)) else { + return false; + }; + let Some(result_path) = 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 cells.closed_cell(ctx, operation.get_result(0)).is_some(); + } + if opid == MirCondBranchOp::get_opid_static() && operand_index == 0 { + return false; + } + mma_branch_forwards_fact(ctx, facts, user, value, 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 { + let scope: HashSet> = operations.iter().copied().collect(); + let cells = LocalCells::new(ctx, operations); + for (value, paths) in &facts.facts { + for (path, found) in paths { + for fact in found.iter().copied() { + if fact.is_opaque() { + continue; + } + for r#use in value.uses(ctx) { + let user = r#use.user_op(); + let operand_index = r#use.find_index(ctx); + if scope.contains(&user) + && (mma_semantic_consumer(ctx, user, operand_index, fact) + || mma_ordinary_use_forwards_fact( + ctx, + &cells, + facts, + *value, + path, + 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(); + // Facts start only at this family's operations; a scope without one has nothing to check. + if !scope_has_any(ctx, all_ops, &ids) { + return Ok(()); + } + 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..3787024c30 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/pipeline.rs @@ -0,0 +1,623 @@ +/* + * 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, scope_has_any}, + ghosts::pipeline_semantic_ids, + invalid, + provenance::{ + CarrierFact, LocalCells, 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), +} + +impl CarrierFact for PipelineFact { + fn is_opaque(&self) -> bool { + matches!(self, Self::Unknown(_)) + } +} + +/// 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 cells = LocalCells::new(ctx, operations); + 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, &cells, &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, + &cells, + &mut facts, + PipelineFact::Unknown, + ); + let propagated = propagate_ordinary_plumbing(ctx, operations, &cells, &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(); + // Facts start only at this family's operations; a scope without one has nothing to check. + if !scope_has_any(ctx, all_ops, &ids) { + return Ok(()); + } + 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..3732e626c4 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/provenance.rs @@ -0,0 +1,576 @@ +/* + * 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::{ + basic_block::BasicBlock, + builtin::{op_interfaces::BranchOpInterface, types::IntegerType}, + context::{Context, Ptr}, + op::{Op, op_cast}, + operation::Operation, + r#type::{TypeHandle, Typed}, + value::Value, +}; + +/// One possible origin of an ordinary carrier. +pub(super) trait CarrierFact: Clone + Eq + std::hash::Hash { + /// True for a fact that names an ordinary SSA value instead of a semantic + /// origin: an input without provenance, or an ordinary scalar. + fn is_opaque(&self) -> bool; +} + +/// Consumers read an origin set only through its size, its sole origin, and +/// whether it holds a given semantic origin. Two opaque origins therefore show +/// a merge exactly as any larger number does, and keeping more would let the +/// sets grow with every merge they flow through. +const OPAQUE_FACTS_KEPT: usize = 2; + +/// Monotone sets of possible semantic origins, indexed by carrier and field path. +#[derive(Clone)] +pub(super) struct OrdinaryFacts { + /// Accumulated origins by SSA carrier, then by field path inside it. An + /// empty path means the value itself; a non-empty path describes a field + /// inside an ordinary aggregate before a matching `mir.extract_field` + /// exposes it. Multiple origins at one path represent a merge that + /// consumers must resolve. + pub(super) facts: HashMap, HashSet>>, +} + +impl Default for OrdinaryFacts { + fn default() -> Self { + Self { + facts: HashMap::new(), + } + } +} + +impl OrdinaryFacts { + /// 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 { + let found = self + .facts + .entry(value) + .or_default() + .entry(path) + .or_default(); + if fact.is_opaque() + && found.iter().filter(|kept| kept.is_opaque()).count() >= OPAQUE_FACTS_KEPT + { + return false; + } + found.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(&value)?.get(path) + } + + /// 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 + .get(&value) + .into_iter() + .flatten() + .flat_map(|(path, facts)| facts.iter().cloned().map(|fact| (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 +} + +/// Closedness of every local cell allocated in one scope. The verifier never +/// mutates the IR it reads, so each cell's alias graph is walked once and +/// shared by all accesses and fixed-point iterations. +pub(super) struct LocalCells { + closed: HashMap, +} + +impl LocalCells { + /// Check the alias graph of each `mir.alloca` result in the scope. + pub(super) fn new(ctx: &Context, operations: &[Ptr]) -> Self { + let closed = operations + .iter() + .filter(|operation| { + Operation::get_opid(**operation, ctx) == MirAllocaOp::get_opid_static() + && operation.deref(ctx).get_num_results() == 1 + }) + .map(|operation| { + let root = operation.deref(ctx).get_result(0); + (root, local_cell_is_closed(ctx, root)) + }) + .collect(); + Self { closed } + } + + /// Resolve a pointer to a closed local cell and the field path inside it. + pub(super) fn closed_cell(&self, ctx: &Context, pointer: Value) -> Option<(Value, Vec)> { + let (root, path) = local_cell(ctx, pointer)?; + let closed = self + .closed + .get(&root) + .copied() + .unwrap_or_else(|| local_cell_is_closed(ctx, root)); + closed.then_some((root, path)) + } +} + +/// Read the operands each incoming edge passes to a block, once per edge. +fn incoming_edge_operands(ctx: &Context, block: Ptr) -> Vec> { + block + .uses(ctx) + .into_iter() + .filter_map(|edge| branch_successor_operands(ctx, edge.user_op(), edge.find_index(ctx))) + .collect() +} + +/// 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], + cells: &LocalCells, + facts: &mut OrdinaryFacts, +) -> bool +where + F: CarrierFact, +{ + 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)) = cells.closed_cell(ctx, store.address_opd(ctx)) + { + 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)) = cells.closed_cell(ctx, load.address_opd(ctx)) + { + changed |= facts.copy_entries(root, operation_ref.get_result(0), &path, &[]); + } + } + } + + for block in blocks { + let edges = incoming_edge_operands(ctx, block); + 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 operands in &edges { + 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: CarrierFact, +{ + 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], + cells: &LocalCells, + facts: &mut OrdinaryFacts, + unknown: impl Fn(Value) -> F + Copy, +) -> bool +where + F: CarrierFact, +{ + 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)) = cells.closed_cell(ctx, store.address_opd(ctx)) else { + 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 edges = incoming_edge_operands(ctx, block); + 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 operands in &edges { + 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..205e60a15d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/scheduler.rs @@ -0,0 +1,398 @@ +/* + * 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, scope_has_any}, + ghosts::scheduler_semantic_ids, + invalid, + provenance::{ + CarrierFact, LocalCells, 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), +} + +impl CarrierFact for SchedulerFact { + fn is_opaque(&self) -> bool { + matches!(self, Self::Unknown(_)) + } +} + +/// 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 cells = LocalCells::new(ctx, operations); + 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, &cells, &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, + &cells, + &mut facts, + SchedulerFact::Unknown, + ); + let propagated = propagate_ordinary_plumbing(ctx, operations, &cells, &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(); + // Facts start only at this family's operations; a scope without one has nothing to check. + if !scope_has_any(ctx, all_ops, &ids) { + return Ok(()); + } + 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..4df2c3f17e --- /dev/null +++ b/cuda-oxide/cute-rs-preview/dialect-cute/src/verify/tests.rs @@ -0,0 +1,958 @@ +/* + * 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::{ + CarrierFact, LocalCells, OrdinaryFacts, 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. +/// The scheduler state reaches the loop header through `pad` ordinary blocks, each passing it +/// as block arguments and through a closed local cell. +fn append_complete_scheduler_function( + ctx: &mut Context, + module_block: Ptr, + name: &str, + grid: CuteTileGridAttr, + misroute_work_edge: bool, + pad: usize, +) { + 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); + let mut carried = vec![start.current(ctx), start.stride(ctx)]; + let mut previous = entry; + for _ in 0..pad { + let block = BasicBlock::new(ctx, None, vec![u64_ty, u64_ty]); + block.insert_at_back(region, ctx); + goto(ctx, previous, block, carried); + let current = block.deref(ctx).get_argument(0); + let stride = block.deref(ctx).get_argument(1); + let cell = local_cell(ctx, block, u64_ty); + store(ctx, block, cell, current); + carried = vec![load(ctx, block, cell, u64_ty), stride]; + previous = block; + } + goto(ctx, previous, header, carried); + + 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, + 0, + ); + append_complete_scheduler_function( + &mut ctx, + block, + "second", + CuteTileGridAttr::new(5, 4, 1), + false, + 0, + ); + 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, + 0, + ); + + 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}"); +} + +/// A scope's cell table answers every pointer as walking that pointer's cell aliases does, +/// including cells allocated outside the scope. +#[test] +fn local_cells_agree_with_the_alias_walk() { + 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 closed = local_cell(&mut ctx, block, u64_ty); + let read_only_type: TypeHandle = MirPtrType::get_generic(&mut ctx, u64_ty, false).into(); + let alias_cast = pointer_cast(&mut ctx, block, closed, read_only_type); + let alias = alias_cast.deref(&ctx).get_result(0); + let value = undef(&mut ctx, block, u64_ty); + store(&mut ctx, block, closed, value); + load(&mut ctx, block, alias, u64_ty); + let escaped = local_cell(&mut ctx, block, u64_ty); + let pointer_type = escaped.get_type(&ctx); + let holder = local_cell(&mut ctx, block, pointer_type); + store(&mut ctx, block, holder, escaped); + let operations: Vec<_> = block.deref(&ctx).iter(&ctx).collect(); + let without_allocas: Vec<_> = operations + .iter() + .copied() + .filter(|operation| Operation::get_opid(*operation, &ctx) != MirAllocaOp::get_opid_static()) + .collect(); + + for scope in [&operations, &without_allocas] { + let cells = LocalCells::new(&ctx, scope); + for pointer in [closed, alias, escaped, holder] { + let walked = super::provenance::local_cell(&ctx, pointer) + .filter(|(root, _)| local_cell_is_closed(&ctx, *root)); + assert_eq!(cells.closed_cell(&ctx, pointer), walked); + } + assert!(cells.closed_cell(&ctx, alias).is_some()); + assert!(cells.closed_cell(&ctx, escaped).is_none()); + } +} + +/// Fact lookups see only the named carrier, and a copy forwards the source as it was when the copy began. +#[test] +fn ordinary_facts_forward_a_snapshot_of_one_carrier() { + let mut ctx = Context::new(); + let (_, block) = module_top(&mut ctx); + let u64_ty: TypeHandle = IntegerType::get(&ctx, 64, Signedness::Unsigned).into(); + let source = undef(&mut ctx, block, u64_ty); + let other = undef(&mut ctx, block, u64_ty); + let destination = undef(&mut ctx, block, u64_ty); + let sorted = |facts: &OrdinaryFacts, value| { + let mut entries = facts.entries(value); + entries.sort(); + entries + }; + + let mut facts = OrdinaryFacts::default(); + assert!(facts.insert_path(source, vec![0], 1)); + assert!(facts.insert_path(source, vec![1, 2], 2)); + assert!(!facts.insert_path(source, vec![1, 2], 2)); + assert!(facts.insert(other, 3)); + assert_eq!(sorted(&facts, source), vec![(vec![0], 1), (vec![1, 2], 2)]); + assert_eq!(sorted(&facts, other), vec![(vec![], 3)]); + assert!(sorted(&facts, destination).is_empty()); + assert_eq!( + facts.at_path(source, &[1, 2]).map(|found| found.len()), + Some(1) + ); + assert!(facts.at_path(source, &[1]).is_none()); + assert!(facts.at(other).is_some_and(|found| found.contains(&3))); + + assert!(facts.copy_entries(source, source, &[], &[5])); + assert_eq!( + sorted(&facts, source), + vec![ + (vec![0], 1), + (vec![1, 2], 2), + (vec![5, 0], 1), + (vec![5, 1, 2], 2) + ] + ); + assert!(facts.copy_entries(source, destination, &[1], &[])); + assert!(!facts.copy_entries(source, destination, &[1], &[])); + assert_eq!(sorted(&facts, destination), vec![(vec![2], 2)]); + assert_eq!(sorted(&facts, other), vec![(vec![], 3)]); + + // Opaque origins stop at two per path; semantic origins are always kept. + assert!(facts.insert(other, 100)); + assert!(facts.insert(other, 101)); + assert!(!facts.insert(other, 102)); + assert!(!facts.insert(other, 100)); + assert!(facts.insert(other, 4)); + assert_eq!( + sorted(&facts, other), + vec![(vec![], 3), (vec![], 4), (vec![], 100), (vec![], 101)] + ); + assert!(facts.insert(destination, 102)); + assert!(facts.copy_entries(destination, other, &[], &[])); + assert_eq!( + sorted(&facts, other), + vec![ + (vec![], 3), + (vec![], 4), + (vec![], 100), + (vec![], 101), + (vec![2], 2) + ] + ); +} + +/// Origins at or above 100 stand for opaque facts in the fact-set tests. +impl CarrierFact for u32 { + fn is_opaque(&self) -> bool { + *self >= 100 + } +} + +/// A scheduler cursor that passes through many blocks and cells keeps its one origin. +#[test] +fn accepts_scheduler_state_carried_through_many_blocks_and_cells() { + let mut ctx = Context::new(); + let (module, block) = module_top(&mut ctx); + append_complete_scheduler_function( + &mut ctx, + block, + "padded", + CuteTileGridAttr::new(4, 4, 1), + false, + 48, + ); + verify_cute_semantics(&ctx, module).unwrap(); +} + +/// Build a function whose `cute.scheduler_has_work` reads a cursor with no scheduler origin: +/// a plain value, or a block argument merging `incoming` unrelated values. +fn scheduler_has_work_without_origin(incoming: usize) -> String { + 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, + ); + let function_op = MirFuncOp::new(&mut ctx, function, TypeAttr::new(function_type.into())); + function_op.set_symbol_name(&mut ctx, "unrooted".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(&mut ctx, None, vec![]); + entry.insert_at_back(region, &ctx); + let mut cursor = undef(&mut ctx, entry, u64_ty); + let mut current = entry; + if incoming > 0 { + let join = BasicBlock::new(&mut ctx, None, vec![u64_ty]); + join.insert_at_back(region, &ctx); + goto(&mut ctx, entry, join, vec![cursor]); + for _ in 1..incoming { + let side = BasicBlock::new(&mut ctx, None, vec![]); + side.insert_at_back(region, &ctx); + let value = undef(&mut ctx, side, u64_ty); + goto(&mut ctx, side, join, vec![value]); + } + cursor = join.deref(&ctx).get_argument(0); + current = join; + } + let grid = CuteTileGridAttr::new(4, 4, 1); + CuteSchedulerHasWorkOp::new(&mut ctx, cursor, grid) + .get_operation() + .insert_at_back(current, &ctx); + verify_cute_semantics(&ctx, module).unwrap_err().to_string() +} + +/// Rejections keep their wording however many unrelated values a cursor merges. +#[test] +fn scheduler_cursor_without_origin_is_rejected_with_the_same_wording() { + let prefix = "invalid CuTe semantic module: cute.scheduler_has_work current"; + assert_eq!( + scheduler_has_work_without_origin(0), + format!("{prefix} is not derived from the required scheduler_new_1d result") + ); + for incoming in [1, 3, 8] { + assert_eq!( + scheduler_has_work_without_origin(incoming), + format!("{prefix} merges values from different scheduler stories") + ); + } +} + +/// 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..b53df69613 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_ctm/Cargo.lock @@ -0,0 +1,765 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..64ee1118e4 --- /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(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT)) + .expect("create start event"); + let end = ctx + .new_event(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..d1c4551376 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/blockscale_gemm_cute/Cargo.lock @@ -0,0 +1,777 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..72b67ce03f --- /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(Some(cuda_core::sys::CUevent_flags_enum_CU_EVENT_DEFAULT)) + .expect("create start event"); + let end = ctx + .new_event(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..f9407efb1d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/elementwise_cute/Cargo.lock @@ -0,0 +1,776 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..967b24730a --- /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::simt::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..1bcc8c3ca2 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/ffn_cute/Cargo.lock @@ -0,0 +1,777 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..ba82a145c8 --- /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::simt::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..fd61d454c1 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_decode_sm120a_cute/Cargo.lock @@ -0,0 +1,777 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..ea8e987fb4 --- /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(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + ctx.new_event(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..5cf93cc223 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fmha_prefill_sm120a_cute/Cargo.lock @@ -0,0 +1,777 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..e26ac291e9 --- /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(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + ctx.new_event(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..c1474dc66d --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/fp16_gemm_256x352_ctm/Cargo.lock @@ -0,0 +1,765 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..79aa0e5756 --- /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(Some(sys::CUevent_flags_enum_CU_EVENT_DEFAULT))?, + ctx.new_event(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..8b256732c0 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/nvfp4_gemv_cute/Cargo.lock @@ -0,0 +1,776 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..f81772a990 --- /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::simt::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..89b1344256 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rmsnorm_cute/Cargo.lock @@ -0,0 +1,777 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..8c79daf8a6 --- /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::simt::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..451bd568f8 --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/rope_cache_cute/Cargo.lock @@ -0,0 +1,777 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..262371d2e2 --- /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::simt::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..ac93a53faa --- /dev/null +++ b/cuda-oxide/cute-rs-preview/examples/sm100_fmha_prefill_cute/Cargo.lock @@ -0,0 +1,778 @@ +# 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 = "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", + "cuda-bindings", + "cuda-core-derive", + "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 = [ + "cuda-macros", +] + +[[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-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 = [ + "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..8a1d5cdb33 --- /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::simt::LaunchConfig { + cuda_core::simt::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::simt::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..46bf88a5fd --- /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::simt::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::